AVX512 for PQFastScan

Signed-off-by: Alexandr Guzhva <[email protected]>
facebookresearch · Mar 18, 2024 · 182affc · 182affc
1 parent 7d21c92
commit 182affc
Show file tree

Hide file tree

Showing 5 changed files with 784 additions and 2 deletions.
diff --git a/faiss/impl/LookupTableScaler.h b/faiss/impl/LookupTableScaler.h
@@ -38,6 +38,23 @@ struct DummyScaler {
         return simd16uint16(0);
     }
 
+#ifdef __AVX512F__
+    inline simd64uint8 lookup(const simd64uint8&, const simd64uint8&) const {
+        FAISS_THROW_MSG("DummyScaler::lookup should not be called.");
+        return simd64uint8(0);
+    }
+
+    inline simd32uint16 scale_lo(const simd64uint8&) const {
+        FAISS_THROW_MSG("DummyScaler::scale_lo should not be called.");
+        return simd32uint16(0);
+    }
+
+    inline simd32uint16 scale_hi(const simd64uint8&) const {
+        FAISS_THROW_MSG("DummyScaler::scale_hi should not be called.");
+        return simd32uint16(0);
+    }
+#endif
+
     template <class dist_t>
     inline dist_t scale_one(const dist_t&) const {
         FAISS_THROW_MSG("DummyScaler::scale_one should not be called.");
@@ -67,6 +84,23 @@ struct NormTableScaler {
         return (simd16uint16(res) >> 8) * scale_simd;
     }
 
+#ifdef __AVX512F__
+    inline simd64uint8 lookup(const simd64uint8& lut, const simd64uint8& c)
+            const {
+        return lut.lookup_4_lanes(c);
+    }
+
+    inline simd32uint16 scale_lo(const simd64uint8& res) const {
+        auto scale_simd_wide = simd32uint16(scale_simd, scale_simd);
+        return simd32uint16(res) * scale_simd_wide;
+    }
+
+    inline simd32uint16 scale_hi(const simd64uint8& res) const {
+        auto scale_simd_wide = simd32uint16(scale_simd, scale_simd);
+        return (simd32uint16(res) >> 8) * scale_simd_wide;
+    }
+#endif
+
     // for non-SIMD implem 2, 3, 4
     template <class dist_t>
     inline dist_t scale_one(const dist_t& x) const {