Enable Miri tests in CI with non-blocking execution

.github/workflows/ci.yml

@@ -204,28 +204,30 @@ jobs:
   #       env:
   #         CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
 
-  # miri:
-  #   needs: basics
-  #   name: miri-test
-  #   runs-on: ubuntu-latest
-  #   steps:
-  #     - uses: actions/checkout@v4
-  #       with:
-  #         lfs: true
+  miri:
+    needs: basics
+    name: miri-test
+    runs-on: ubuntu-latest
+    continue-on-error: true
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          lfs: true
 
-  #     - name: Install Rust nightly with miri
-  #       uses: dtolnay/rust-toolchain@stable
-  #       with:
-  #         toolchain: nightly
-  #         components: miri
+      - name: Install Rust nightly with miri
+        uses: dtolnay/rust-toolchain@stable
+        with:
+          toolchain: nightly
+          components: miri
 
-  #     - name: Install cargo-nextest
-  #       uses: taiki-e/install-action@v2
-  #       with:
-  #         tool: cargo-nextest
+      - name: Install cargo-nextest
+        uses: taiki-e/install-action@v2
+        with:
+          tool: cargo-nextest
 
-  #     - uses: Swatinem/rust-cache@v2
-  #     - name: miri
-  #       run: cargo +nightly miri nextest run --package diskann-quantization
-  #       env:
-  #         MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance
+      - uses: Swatinem/rust-cache@v2
+
+      - name: miri
+        run: cargo +nightly miri nextest run --package diskann-quantization
+        env:
+          MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance

diskann-quantization/src/algorithms/heap.rs

@@ -390,11 +390,15 @@ mod tests {
         // Heap of size 2.
         fuzz_test_impl(2, 101, &mut rng);
 
-        // Heap size not power of two.
-        fuzz_test_impl(1000, 1000, &mut rng);
-
-        // Heap size power of two.
-        fuzz_test_impl(128, 1000, &mut rng);
+        // Miri is extremely slow, so we skip the larger tests there.
+        #[cfg(not(miri))]
+        {
+            // Heap size not power of two.
+            fuzz_test_impl(1000, 1000, &mut rng);
+
+            // Heap size power of two.
+            fuzz_test_impl(128, 1000, &mut rng);
+        }
     }
 
     #[test]

diskann-quantization/src/algorithms/kmeans/common.rs

@@ -565,19 +565,35 @@ mod tests {
 
     #[test]
     fn test_block_transpose_16() {
+        #[cfg(not(miri))]
         for nrows in 0..128 {
             for ncols in 0..5 {
                 test_block_transpose::<16>(nrows, ncols);
             }
         }
+
+        #[cfg(miri)]
+        for nrows in 127..128 {
+            for ncols in 4..5 {
+                test_block_transpose::<16>(nrows, ncols);
+            }
+        }
     }
 
     #[test]
     fn test_block_transpose_8() {
+        #[cfg(not(miri))]
         for nrows in 0..128 {
             for ncols in 0..5 {
                 test_block_transpose::<8>(nrows, ncols);
             }
         }
+
+        #[cfg(miri)]
+        for nrows in 127..128 {
+            for ncols in 4..5 {
+                test_block_transpose::<8>(nrows, ncols);
+            }
+        }
     }
 }

diskann-quantization/src/algorithms/kmeans/lloyds.rs

@@ -457,14 +457,15 @@ pub fn lloyds(
 
 #[cfg(test)]
 mod tests {
-    use diskann_utils::{lazy_format, views::Matrix};
+    #[cfg(not(miri))]
+    use diskann_utils::lazy_format;
+    use diskann_utils::views::Matrix;
     use diskann_vector::{distance::SquaredL2, PureDistanceFunction};
-    use rand::{
-        distr::{Distribution, Uniform},
-        rngs::StdRng,
-        seq::{IndexedRandom, SliceRandom},
-        Rng, SeedableRng,
-    };
+    #[cfg(not(miri))]
+    use rand::distr::{Distribution, Uniform};
+    #[cfg(not(miri))]
+    use rand::seq::IndexedRandom;
+    use rand::{rngs::StdRng, seq::SliceRandom, Rng, SeedableRng};
 
     use super::*;
 
@@ -477,6 +478,7 @@ mod tests {
     // relatively quickly.
     //
     // Outside of rare validations, Miri tests go through a different path for speed purposes.
+    #[cfg(not(miri))]
     fn test_distances_in_place_impl<R: Rng>(
         ndata: usize,
         ncenters: usize,
@@ -556,15 +558,11 @@ mod tests {
         }
     }
 
-    cfg_if::cfg_if! {
-        if #[cfg(miri)] {
-            const TRIALS: usize = 1;
-        } else {
-            const TRIALS: usize = 100;
-        }
-    }
+    #[cfg(not(miri))]
+    const TRIALS: usize = 100;
 
     #[test]
+    #[cfg(not(miri))]
     fn test_distances_in_place() {
         let mut rng = StdRng::seed_from_u64(0xece88a9c6cd86a8a);
         for ndata in 1..=31 {
@@ -719,12 +717,22 @@ mod tests {
     #[test]
     fn end_to_end_test() {
         let mut rng = StdRng::seed_from_u64(0xff22c38d0f0531bf);
-        let setup = EndToEndSetup {
-            ncenters: 11,
-            ndim: 4,
-            data_per_center: 8,
-            step_between_clusters: 20,
-            ntrials: 10,
+        let setup = if cfg!(miri) {
+            EndToEndSetup {
+                ncenters: 3,
+                ndim: 4,
+                data_per_center: 2,
+                step_between_clusters: 20,
+                ntrials: 2,
+            }
+        } else {
+            EndToEndSetup {
+                ncenters: 11,
+                ndim: 4,
+                data_per_center: 8,
+                step_between_clusters: 20,
+                ntrials: 10,
+            }
         };
         end_to_end_test_impl(&setup, &mut rng);
     }

diskann-quantization/src/algorithms/kmeans/plusplus.rs

@@ -683,6 +683,11 @@ mod tests {
     fn test_update_distances() {
         let mut rng = StdRng::seed_from_u64(0x56c94b53c73e4fd9);
         for num_points in 0..48 {
+            #[cfg(miri)]
+            if num_points % 7 != 0 {
+                continue;
+            }
+
             for dim in 1..4 {
                 test_update_distances_impl(num_points, dim, &mut rng);
             }
@@ -695,6 +700,7 @@ mod tests {
 
     // Kmeans++ sanity checks - if there are only `N` distinct and we want `N` centers,
     // then all `N` should be selected without repeats.
+    #[cfg(not(miri))]
     fn sanity_check_impl<R: Rng>(ncenters: usize, dim: usize, rng: &mut R) {
         let repeats_per_center = 3;
         let context = lazy_format!(
@@ -756,6 +762,7 @@ mod tests {
 
     // This test is like the sanity check - but instead of exact repeats, we use slightly
     // perturbed values to test that the proportionality is of distances is respected.
+    #[cfg(not(miri))]
     fn fuzzy_sanity_check_impl<R: Rng>(ncenters: usize, dim: usize, rng: &mut R) {
         let repeats_per_center = 3;
 

diskann-quantization/src/algorithms/transforms/double_hadamard.rs

@@ -387,6 +387,7 @@ where
 ///////////
 
 #[cfg(test)]
+#[cfg(not(miri))]
 mod tests {
     use diskann_utils::lazy_format;
     use rand::{rngs::StdRng, SeedableRng};

diskann-quantization/src/algorithms/transforms/mod.rs

@@ -28,6 +28,7 @@ crate::utils::features! {
 mod utils;
 
 #[cfg(test)]
+#[cfg(not(miri))]
 mod test_utils;
 
 // reexports
@@ -355,4 +356,5 @@ pub enum TargetDim {
 }
 
 #[cfg(test)]
+#[cfg(not(miri))]
 test_utils::delegate_transformer!(Transform<crate::alloc::GlobalAllocator>);

diskann-quantization/src/algorithms/transforms/padding_hadamard.rs

@@ -366,14 +366,15 @@ where
 
 #[cfg(test)]
 mod tests {
+    #[cfg(not(miri))]
     use diskann_utils::lazy_format;
     use rand::{rngs::StdRng, SeedableRng};
 
     use super::*;
-    use crate::{
-        algorithms::transforms::{test_utils, Transform, TransformKind},
-        alloc::GlobalAllocator,
-    };
+
+    #[cfg(not(miri))]
+    use crate::algorithms::transforms::{test_utils, Transform, TransformKind};
+    use crate::alloc::GlobalAllocator;
 
     // Since we use a slightly non-obvious strategy for applying the +/-1 permutation, we
     // test its behavior explicitly.
@@ -441,11 +442,13 @@ mod tests {
         assert_eq!(output[15], 0.0f32);
     }
 
+    #[cfg(not(miri))]
     test_utils::delegate_transformer!(PaddingHadamard<GlobalAllocator>);
 
     // This tests the natural hadamard transform where the output dimension is upgraded
     // to the next power of 2.
     #[test]
+    #[cfg(not(miri))]
     fn test_padding_hadamard() {
         // Inner product computations are more susceptible to floating point error.
         // Instead of using ULP here, we fall back to using absolute and relative error.

diskann-quantization/src/algorithms/transforms/test_utils.rs

@@ -172,7 +172,11 @@ fn within_ulp(mut got: f32, expected: f32, ulp: usize) -> bool {
 #[derive(Debug, Clone, Copy)]
 pub(super) enum Check {
     Ulp(usize),
-    AbsRel { abs: f32, rel: f32 },
+    AbsRel {
+        abs: f32,
+        rel: f32,
+    },
+    #[cfg(not(miri))]
     Skip,
 }
 
@@ -185,6 +189,7 @@ impl Check {
         Self::AbsRel { abs, rel }
     }
 
+    #[cfg(not(miri))]
     pub(super) fn skip() -> Self {
         Self::Skip
     }
@@ -219,6 +224,7 @@ impl Check {
                     })
                 }
             }
+            #[cfg(not(miri))]
             Self::Skip => Ok(()),
         }
     }

diskann-quantization/src/bits/distances.rs

@@ -2003,6 +2003,12 @@ mod tests {
         let dist = Uniform::new_inclusive(min, max).unwrap();
 
         for dim in 0..dim_max {
+            // Only run the maximum dimension when running under miri.
+            #[cfg(miri)]
+            if dim != dim_max - 1 {
+                continue;
+            }
+
             let mut x_reference: Vec<u8> = vec![0; dim];
             let mut y_reference: Vec<u8> = vec![0; dim];
 
@@ -2092,7 +2098,7 @@ mod tests {
 
     cfg_if::cfg_if! {
         if #[cfg(miri)] {
-            const MAX_DIM: usize = 128;
+            const MAX_DIM: usize = 8;
             const TRIALS_PER_DIM: usize = 1;
         } else {
             const MAX_DIM: usize = 256;

diskann-quantization/src/bits/slice.rs

@@ -1512,6 +1512,11 @@ mod tests {
     fn test_binary_dense() {
         let mut rng = StdRng::seed_from_u64(0xb3c95e8e19d3842e);
         for len in 0..MAX_DIM {
+            #[cfg(miri)]
+            if len != MAX_DIM - 1 {
+                continue;
+            }
+
             test_send_and_sync::<1, Binary, Dense>();
             test_empty::<1, Binary, Dense>();
             test_construction_errors::<1, Binary, Dense>();
@@ -1558,6 +1563,11 @@ mod tests {
     fn test_4bit_bit_transpose() {
         let mut rng = StdRng::seed_from_u64(0xb3c95e8e19d3842e);
         for len in 0..MAX_DIM {
+            #[cfg(miri)]
+            if len != MAX_DIM - 1 {
+                continue;
+            }
+
             test_send_and_sync::<4, Unsigned, BitTranspose>();
             test_empty::<4, Unsigned, BitTranspose>();
             test_construction_errors::<4, Unsigned, BitTranspose>();

diskann-quantization/src/minmax/quantizer.rs

@@ -337,6 +337,7 @@ impl_functor!(MinMaxCosineNormalized);
 // Tests //
 ///////////
 #[cfg(test)]
+#[cfg(not(miri))]
 mod minmax_quantizer_tests {
     use std::num::NonZeroUsize;
 

diskann-quantization/src/minmax/vectors.rs

@@ -489,6 +489,7 @@ where
 ///////////
 
 #[cfg(test)]
+#[cfg(not(miri))]
 mod minmax_vector_tests {
     use diskann_utils::Reborrow;
     use rand::{
@@ -752,15 +753,16 @@ mod minmax_vector_tests {
             #[test]
             fn $name() {
                 let mut rng = StdRng::seed_from_u64($seed);
-                for dim in 1..(bit_scale::<$nbits>() as usize) {
+                const MAX_DIM: usize = (bit_scale::<$nbits>() as usize);
+                for dim in 1..=MAX_DIM {
                     for _ in 0..TRIALS {
                         test_minmax_compensated_vectors::<$nbits, _>(dim, &mut rng);
                     }
                 }
             }
         };
     }
-    test_minmax_compensated!(unsigned_minmax_compensated_test_u1, 1, 0xa32d5658097a1c35);
+    test_minmax_compensated!(unsigned_minmax_compensated_test_u1, 1, 0xa33d5658097a1c35);
     test_minmax_compensated!(unsigned_minmax_compensated_test_u2, 2, 0xaedf3d2a223b7b77);
     test_minmax_compensated!(unsigned_minmax_compensated_test_u4, 4, 0xf60c0c8d1aadc126);
     test_minmax_compensated!(unsigned_minmax_compensated_test_u8, 8, 0x09fa14c42a9d7d98);