diff --git a/.jules/bolt.md b/.jules/bolt.md index 19a1db4..a61b336 100644 --- a/.jules/bolt.md +++ b/.jules/bolt.md @@ -1,3 +1,6 @@ ## 2024-05-18 - Fast row-wise Euclidean norm in pure NumPy **Learning:** In performance-critical paths, computing the batch norm of a 2D array via `np.linalg.norm(arr, axis=1)` is relatively slow. Using `np.sqrt(np.einsum('ij,ij->i', arr, arr))` is significantly faster (~4x speedup on a laptop CPU for typical batch sizes). If `keepdims=True` behavior is needed, appending `[:, np.newaxis]` matches the original shape seamlessly. **Action:** Always prefer `np.sqrt(np.einsum('ij,ij->i', arr, arr))` over `np.linalg.norm(arr, axis=1)` when computing row-wise vector norms in NumPy to eliminate dispatch overhead and improve execution speed. +## 2025-02-18 - Fast batched squared norms for 3D arrays and Contiguous Matrix Multiplications +**Learning:** Using `np.einsum('ijk,ijk->ij', X, X)` is significantly faster than `(X ** 2).sum(2)` for calculating squared Euclidean norms along the last axis of a 3D NumPy array, avoiding large intermediate array allocations. Also, explicitly using associativity (e.g. `R @ S.T` instead of `(S @ R.T).T`) yields a C-contiguous array instead of an F-contiguous view, which speeds up the operation and subsequent steps relying on cache locality. +**Action:** Replace `(X ** 2).sum(2)` with `np.einsum('ijk,ijk->ij', X, X)` for 3D array batched squared norm calculations and rewrite expression to avoid explicit transpositions to preserve C-contiguity in critical data paths. diff --git a/snapvec/_index.py b/snapvec/_index.py index fdc793e..2696754 100644 --- a/snapvec/_index.py +++ b/snapvec/_index.py @@ -303,7 +303,9 @@ def add_batch(self, ids: list[Any], vectors: NDArray[np.float32]) -> None: reconstructed: NDArray[np.float32] = self._centroids[batch_idx] r_scaled: NDArray[np.float32] = scaled - reconstructed # sign(S·r_rot) = sign(S·r_scaled) — scale-invariant - S_r: NDArray[np.float32] = (self._S @ r_scaled.T).T + # Optimized: Avoid explicit transpositions and intermediate allocations. + # Yields a C-contiguous array rather than an F-contiguous view. + S_r: NDArray[np.float32] = r_scaled @ self._S.T qjl_signs = np.sign(S_r).astype(np.int8) qjl_signs[qjl_signs == 0] = 1 # Store ‖r_rot‖ = ‖r_scaled‖/√pdim (unscaled space norm) diff --git a/snapvec/_ivfpq.py b/snapvec/_ivfpq.py index bcf3e51..81f6422 100644 --- a/snapvec/_ivfpq.py +++ b/snapvec/_ivfpq.py @@ -429,7 +429,8 @@ def add_batch( if self.keep_full_precision else np.empty((0, self._pdim), dtype=np.float16) ) - cb_norms = (self._codebooks ** 2).sum(2) # (M, K) + # Optimized: np.einsum('ijk,ijk->ij', X, X) is faster and avoids large intermediate arrays + cb_norms = np.einsum('ijk,ijk->ij', self._codebooks, self._codebooks) # (M, K) cb_T = np.transpose(self._codebooks, (0, 2, 1)) # (M, d_sub, K) for start in range(0, n, self._ENCODE_CHUNK): end = min(start + self._ENCODE_CHUNK, n)