feat: prewarm (tensorchord#15)

Signed-off-by: usamoi <usamoi@outlook.com>

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"]

src/algorithm/rabitq.rs

@@ -47,6 +47,12 @@ fn orthogonal_matrix(n: usize) -> Vec<Vec<f32>> {
 
 static MATRIXS: [OnceLock<Vec<Vec<f32>>>; 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<f32> {
     use base::scalar::ScalarLike;
     let n = vector.len();

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());
     }
 }

src/gucs/prewarm.rs

@@ -0,0 +1,31 @@
+use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting};
+use std::ffi::CStr;
+
+static PREWARM_DIM: GucSetting<Option<&CStr>> = GucSetting::<Option<&CStr>>::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::<usize>() {
+                    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");
+        }
+    }
+}