diff --git a/README.md b/README.md index 75b4449..7bfe174 100644 --- a/README.md +++ b/README.md @@ -245,6 +245,7 @@ Deprecated datasets will remain available, but their benchmark results will not | [GGNN](https://github.com/cgtuebingen/ggnn) | 0.9 | | [Glass](https://github.com/zilliztech/pyglass) | git+d2296ec | | [HNSW](https://github.com/nmslib/hnswlib) | 0.8.0 | +| [HNSW (Hann)](https://github.com/habedi/hann) | 0.8.3 | | [HNSW-RaBitQ](https://github.com/VectorDB-NTU/RaBitQ-Library) | git+5ea4df0 | | [IVF (Faiss)](https://github.com/facebookresearch/faiss) | 1.14.3 | | [IVF-PQ (Faiss)](https://github.com/facebookresearch/faiss) | 1.14.3 | @@ -260,9 +261,11 @@ Deprecated datasets will remain available, but their benchmark results will not | [NSG](https://github.com/facebookresearch/faiss) | 1.14.3 | | [PAG](https://github.com/KejingLu-810/PAG) | git+ee34ed7 | | [PDX](https://github.com/cwida/PDX) | git+93531b9 | +| [PQIVF (Hann)](https://github.com/habedi/hann) | 0.8.3 | | [PUFFINN](https://github.com/puffinn/puffinn) | git+fd86b0d | | [PyNNDescent](https://github.com/lmcinnes/pynndescent) | 0.6.0 | | [RoarGraph](https://github.com/matchyc/RoarGraph) | git+f2b49b6 | +| [RPT (Hann)](https://github.com/habedi/hann) | 0.8.3 | | [ScaNN](https://github.com/google-research/google-research/tree/master/scann) | 1.4.2 | | [SymphonyQG](https://github.com/gouyt13/SymphonyQG) | git+32a0019 | | [Vamana (DiskANN)](https://github.com/microsoft/DiskANN) | 0.7.0 | diff --git a/install.sh b/install.sh index 6cbee87..7b80e21 100755 --- a/install.sh +++ b/install.sh @@ -73,9 +73,14 @@ build_singularity_image() { if [ ! -e "$image_dir/${name}.sif" ] || [ "$force_build" = "true" ]; then cp "$1/image.def" "$build_images_dir/${name}.def" + # Stage the algorithm directory so the def's %files section can copy + # extra files (sources are resolved relative to the build directory). + rm -rf "$build_images_dir/${name}.files" + cp -r "$1" "$build_images_dir/${name}.files" pushd "$build_images_dir" >/dev/null singularity build -F "${name}.sif" "${name}.def" popd >/dev/null + rm -rf "$build_images_dir/${name}.files" if [ "$build_images_dir" != "$image_dir" ]; then mv "$build_images_dir/${name}.sif" "$image_dir/${name}.sif" fi @@ -99,6 +104,7 @@ clean_up() { ARG=$? rm -f "$build_images_dir/environment.yml" find "$build_images_dir" -maxdepth 1 -name "*.def" -type f -exec rm {} + + find "$build_images_dir" -maxdepth 1 -name "*.files" -type d -exec rm -rf {} + exit $ARG } trap clean_up EXIT diff --git a/vibe/algorithms/hann/bridge.go b/vibe/algorithms/hann/bridge.go new file mode 100644 index 0000000..88221de --- /dev/null +++ b/vibe/algorithms/hann/bridge.go @@ -0,0 +1,305 @@ +// Package main builds a c-shared library that exposes the Hann indexes +// (HNSW, PQIVF, and RPT) through a small C ABI, so the Python wrapper can +// drive them with ctypes. Go objects never cross the boundary: each index +// lives in a mutex-guarded map and is referred to by an int64 handle. +package main + +/* +#include +*/ +import "C" + +import ( + "strings" + "sync" + "unsafe" + + "github.com/habedi/hann/core" + "github.com/habedi/hann/hnsw" + "github.com/habedi/hann/pqivf" + "github.com/habedi/hann/rpt" +) + +type indexEntry struct { + index core.Index + dim int + count int // running count of inserted vectors, used to assign ids +} + +var ( + registry = make(map[int64]*indexEntry) + registryMu sync.Mutex + nextHandle int64 = 1 +) + +func getEntry(handle int64) *indexEntry { + registryMu.Lock() + defer registryMu.Unlock() + return registry[handle] +} + +func putEntry(index core.Index, dim int) C.int64_t { + registryMu.Lock() + defer registryMu.Unlock() + h := nextHandle + nextHandle++ + registry[h] = &indexEntry{index: index, dim: dim} + return C.int64_t(h) +} + +func parseMetric(metric *C.char) (core.Metric, bool) { + switch strings.ToLower(C.GoString(metric)) { + case "euclidean": + return core.Euclidean, true + case "cosine": + return core.Cosine, true + default: + return core.Metric{}, false + } +} + +// hann_hnsw_new creates an HNSW index and returns a handle to it, or -1 on +// error. The metric string is either "euclidean" or "cosine". +// +//export hann_hnsw_new +func hann_hnsw_new(dim, m, efConstruction C.int64_t, metric *C.char) (handle C.int64_t) { + defer func() { + if r := recover(); r != nil { + handle = -1 + } + }() + + coreMetric, ok := parseMetric(metric) + if !ok { + return -1 + } + + index, err := hnsw.New(int(dim), + hnsw.WithM(int(m)), + hnsw.WithEfConstruction(int(efConstruction)), + hnsw.WithMetric(coreMetric), + ) + if err != nil { + return -1 + } + return putEntry(index, int(dim)) +} + +// hann_pqivf_new creates a PQIVF index and returns a handle to it, or -1 on +// error. The metric is always Euclidean, so there is no metric argument. A +// zero for coarseK, numSubquantizers, pqK, kMeansIters, or candidateClusters +// keeps the library default for that parameter. +// +//export hann_pqivf_new +func hann_pqivf_new(dim, coarseK, numSubquantizers, pqK, kMeansIters, candidateClusters C.int64_t) (handle C.int64_t) { + defer func() { + if r := recover(); r != nil { + handle = -1 + } + }() + + var opts []pqivf.Option + if coarseK > 0 { + opts = append(opts, pqivf.WithCoarseK(int(coarseK))) + } + if numSubquantizers > 0 { + opts = append(opts, pqivf.WithNumSubquantizers(int(numSubquantizers))) + } + if pqK > 0 { + opts = append(opts, pqivf.WithPQK(int(pqK))) + } + if kMeansIters > 0 { + opts = append(opts, pqivf.WithKMeansIters(int(kMeansIters))) + } + if candidateClusters > 0 { + opts = append(opts, pqivf.WithCandidateClusters(int(candidateClusters))) + } + + index, err := pqivf.New(int(dim), opts...) + if err != nil { + return -1 + } + return putEntry(index, int(dim)) +} + +// hann_rpt_new creates an RPT index and returns a handle to it, or -1 on +// error. The metric string is either "euclidean" or "cosine". A zero for +// leafCapacity or candidateProjections, and a negative probeMargin, keep the +// library default for that parameter. +// +//export hann_rpt_new +func hann_rpt_new(dim, leafCapacity, candidateProjections C.int64_t, probeMargin C.double, metric *C.char) (handle C.int64_t) { + defer func() { + if r := recover(); r != nil { + handle = -1 + } + }() + + coreMetric, ok := parseMetric(metric) + if !ok { + return -1 + } + + opts := []rpt.Option{rpt.WithMetric(coreMetric)} + if leafCapacity > 0 { + opts = append(opts, rpt.WithLeafCapacity(int(leafCapacity))) + } + if candidateProjections > 0 { + opts = append(opts, rpt.WithCandidateProjections(int(candidateProjections))) + } + if probeMargin >= 0 { + opts = append(opts, rpt.WithProbeMargin(float64(probeMargin))) + } + + index, err := rpt.New(int(dim), opts...) + if err != nil { + return -1 + } + return putEntry(index, int(dim)) +} + +// hann_add_batch adds n vectors of the given dimension, laid out row-major +// in flat, assigning ids sequentially from the running count of the index. +// It returns the number of vectors added, or -1 on error. +// +//export hann_add_batch +func hann_add_batch(handle C.int64_t, flat *C.float, n, dim C.int64_t) (added C.int64_t) { + defer func() { + if r := recover(); r != nil { + added = -1 + } + }() + + entry := getEntry(int64(handle)) + if entry == nil || flat == nil || n <= 0 || int(dim) != entry.dim { + return -1 + } + + rows := int(n) + d := int(dim) + src := unsafe.Slice((*float32)(unsafe.Pointer(flat)), rows*d) + + // Copy every row into a fresh Go slice: the C buffer is owned by the + // caller and must not be retained past this call. + const chunk = 50000 + for start := 0; start < rows; start += chunk { + end := start + chunk + if end > rows { + end = rows + } + batch := make(map[int][]float32, end-start) + for i := start; i < end; i++ { + row := make([]float32, d) + copy(row, src[i*d:(i+1)*d]) + batch[entry.count+i] = row + } + if err := core.BulkAdd(entry.index, batch); err != nil { + return -1 + } + } + entry.count += rows + return n +} + +// hann_train trains the index behind the handle, which must implement +// core.Trainer. It returns 0 on success and -1 on error. +// +//export hann_train +func hann_train(handle C.int64_t) (status C.int64_t) { + defer func() { + if r := recover(); r != nil { + status = -1 + } + }() + + entry := getEntry(int64(handle)) + if entry == nil { + return -1 + } + trainer, ok := entry.index.(core.Trainer) + if !ok { + return -1 + } + if err := trainer.Train(); err != nil { + return -1 + } + return 0 +} + +// hann_hnsw_set_ef changes the search breadth of an HNSW index. It returns 0 +// on success and -1 on error, which includes a handle that does not refer to +// an HNSW index. +// +//export hann_hnsw_set_ef +func hann_hnsw_set_ef(handle, ef C.int64_t) (status C.int64_t) { + defer func() { + if r := recover(); r != nil { + status = -1 + } + }() + + entry := getEntry(int64(handle)) + if entry == nil { + return -1 + } + index, ok := entry.index.(*hnsw.Index) + if !ok { + return -1 + } + if err := index.SetEf(int(ef)); err != nil { + return -1 + } + return 0 +} + +// hann_search searches the index for the k nearest neighbors of the query +// and writes their ids into out, which must have room for k int32 values. +// It returns the number of ids written, or -1 on error. +// +//export hann_search +func hann_search(handle C.int64_t, query *C.float, dim, k C.int64_t, out *C.int32_t) (found C.int64_t) { + defer func() { + if r := recover(); r != nil { + found = -1 + } + }() + + entry := getEntry(int64(handle)) + if entry == nil || query == nil || out == nil || k <= 0 || int(dim) != entry.dim { + return -1 + } + + d := int(dim) + src := unsafe.Slice((*float32)(unsafe.Pointer(query)), d) + q := make([]float32, d) + copy(q, src) + + neighbors, err := entry.index.Search(q, int(k)) + if err != nil { + return -1 + } + + dst := unsafe.Slice((*int32)(unsafe.Pointer(out)), int(k)) + written := 0 + for _, nb := range neighbors { + if written >= int(k) { + break + } + dst[written] = int32(nb.ID) + written++ + } + return C.int64_t(written) +} + +// hann_free releases the index behind the handle. Freeing an unknown handle +// is a no-op. +// +//export hann_free +func hann_free(handle C.int64_t) { + defer func() { _ = recover() }() + registryMu.Lock() + defer registryMu.Unlock() + delete(registry, int64(handle)) +} + +func main() {} diff --git a/vibe/algorithms/hann/config.yml b/vibe/algorithms/hann/config.yml new file mode 100644 index 0000000..ef8f2ef --- /dev/null +++ b/vibe/algorithms/hann/config.yml @@ -0,0 +1,43 @@ +float: + any: + - base_args: ['@metric'] + constructor: HannHNSW + disabled: false + singularity_image: hann + module: vibe.algorithms.hann + name: hnsw(hann) + run_groups: + base: + args: + M: [16, 24, 32, 48] + efConstruction: [200, 400] + query_args: + ef: [10, 20, 40, 80, 120, 200, 400, 600, 800] + - base_args: ['@metric'] + constructor: HannPQIVF + disabled: false + singularity_image: hann + module: vibe.algorithms.hann + name: pqivf(hann) + run_groups: + # The candidate cluster count is a construction-time parameter in Hann, + # so it lives in args and every combination is a separate build. + base: + args: + coarseK: [64, 256, 1024] + candidateClusters: [1, 2, 4, 8, 16, 32] + - base_args: ['@metric'] + constructor: HannRPT + disabled: false + singularity_image: hann + module: vibe.algorithms.hann + name: rpt(hann) + run_groups: + # The probe margin is a construction-time parameter in Hann, so it + # lives in args and every combination is a separate build. It is a + # fraction of the projection value spread at each tree node, which + # needs Hann v0.8.1 or newer (the HANN_REF pinned in image.def). + base: + args: + leafCapacity: [32, 128] + probeMargin: [0.5, 1.0, 1.5, 2.0, 3.0] diff --git a/vibe/algorithms/hann/go.mod b/vibe/algorithms/hann/go.mod new file mode 100644 index 0000000..b1e315e --- /dev/null +++ b/vibe/algorithms/hann/go.mod @@ -0,0 +1,10 @@ +module hann_bridge + +go 1.23.0 + +require github.com/habedi/hann v0.0.0-00010101000000-000000000000 + +// The image.def clones the Hann repository into hann-src next to this file. +// For a local build, point the replace directive at a Hann working copy, or +// place one at hann-src. +replace github.com/habedi/hann => ./hann-src diff --git a/vibe/algorithms/hann/image.def b/vibe/algorithms/hann/image.def new file mode 100644 index 0000000..7b50843 --- /dev/null +++ b/vibe/algorithms/hann/image.def @@ -0,0 +1,38 @@ +Bootstrap: localimage +From: base.sif + +# install.sh stages the algorithm directory as .files next to this def, +# which is where the %files sources below resolve from. + +%files + hann.files/bridge.go /opt/hann/src/bridge.go + hann.files/go.mod /opt/hann/src/go.mod + +%post + # The Hann release to benchmark. Bump this to test a newer version. + HANN_REF=v0.8.3 + GO_VERSION=1.26.5 + + case "$(uname -m)" in + x86_64) go_arch=amd64 ;; + aarch64|arm64) go_arch=arm64 ;; + *) echo "Unsupported architecture: $(uname -m)" >&2; exit 1 ;; + esac + + curl -fL --retry 5 --retry-all-errors -o /tmp/go.tar.gz "https://go.dev/dl/go${GO_VERSION}.linux-${go_arch}.tar.gz" + tar -C /usr/local -xzf /tmp/go.tar.gz + rm /tmp/go.tar.gz + export PATH=/usr/local/go/bin:$PATH + + # Build the c-shared bridge against the pinned Hann release, and install the + # library outside every bind-mounted path (the vibe repository and the home + # directory are mounted at runtime), where module.py expects to find it. + cd /opt/hann/src + git clone https://github.com/habedi/hann hann-src + git -C hann-src checkout "$HANN_REF" + go mod tidy + CGO_ENABLED=1 go build -buildmode=c-shared -o /opt/hann/libhann_bridge.so . + + go clean -modcache + cd / + rm -rf /opt/hann/src /usr/local/go /root/go /root/.cache/go-build diff --git a/vibe/algorithms/hann/module.py b/vibe/algorithms/hann/module.py new file mode 100644 index 0000000..14ea11e --- /dev/null +++ b/vibe/algorithms/hann/module.py @@ -0,0 +1,213 @@ +import ctypes +import os + +import numpy as np + +from ..base.module import BaseANN + +# The bridge library is baked into the image outside any bind-mounted path +# (the VIBE repository and the home directory are mounted by Singularity, so +# a library placed there would be shadowed at runtime). HANN_BRIDGE_LIB +# overrides the path for local testing. +DEFAULT_LIB_PATH = "/opt/hann/libhann_bridge.so" + + +def _load_bridge(): + lib_path = os.environ.get("HANN_BRIDGE_LIB", DEFAULT_LIB_PATH) + lib = ctypes.CDLL(lib_path) + + lib.hann_hnsw_new.argtypes = [ctypes.c_int64, ctypes.c_int64, ctypes.c_int64, ctypes.c_char_p] + lib.hann_hnsw_new.restype = ctypes.c_int64 + + lib.hann_pqivf_new.argtypes = [ctypes.c_int64] * 6 + lib.hann_pqivf_new.restype = ctypes.c_int64 + + lib.hann_rpt_new.argtypes = [ + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_double, + ctypes.c_char_p, + ] + lib.hann_rpt_new.restype = ctypes.c_int64 + + lib.hann_add_batch.argtypes = [ + ctypes.c_int64, + ctypes.POINTER(ctypes.c_float), + ctypes.c_int64, + ctypes.c_int64, + ] + lib.hann_add_batch.restype = ctypes.c_int64 + + lib.hann_train.argtypes = [ctypes.c_int64] + lib.hann_train.restype = ctypes.c_int64 + + lib.hann_hnsw_set_ef.argtypes = [ctypes.c_int64, ctypes.c_int64] + lib.hann_hnsw_set_ef.restype = ctypes.c_int64 + + lib.hann_search.argtypes = [ + ctypes.c_int64, + ctypes.POINTER(ctypes.c_float), + ctypes.c_int64, + ctypes.c_int64, + ctypes.POINTER(ctypes.c_int32), + ] + lib.hann_search.restype = ctypes.c_int64 + + lib.hann_free.argtypes = [ctypes.c_int64] + lib.hann_free.restype = None + + return lib + + +class _HannBase(BaseANN): + """Shared handle management, insertion, and search for the Hann indexes.""" + + def __init__(self): + self.lib = _load_bridge() + self.handle = None + self.dim = None + + def _new_handle(self, dim): + raise NotImplementedError + + def _prepare(self, X): + """Hook for per-index preprocessing of the input matrix.""" + return X + + def fit(self, X): + X = np.ascontiguousarray(self._prepare(X), dtype=np.float32) + n, dim = X.shape + self.dim = dim + self.handle = self._new_handle(dim) + if self.handle < 0: + raise RuntimeError("hann: failed to create the index") + added = self.lib.hann_add_batch( + self.handle, X.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), n, dim + ) + if added != n: + raise RuntimeError("hann: added %d of %d vectors" % (added, n)) + + def query(self, v, n): + q = np.ascontiguousarray(self._prepare(v.reshape(1, -1))[0], dtype=np.float32) + out = np.empty(n, dtype=np.int32) + count = self.lib.hann_search( + self.handle, + q.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + self.dim, + n, + out.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)), + ) + if count < 0: + raise RuntimeError("hann: search failed") + return out[:count] + + def freeIndex(self): + if self.handle is not None: + self.lib.hann_free(self.handle) + self.handle = None + + def done(self): + self.freeIndex() + + +class HannHNSW(_HannBase): + def __init__(self, metric, M, efConstruction): + # "normalized" vectors are pre-normalized, so cosine ordering equals + # inner-product ordering. Plain "ip" has no counterpart in Hann. + if metric not in ("euclidean", "cosine", "normalized"): + raise NotImplementedError(f"HannHNSW does not support metric {metric}") + super().__init__() + self.metric = {"euclidean": "euclidean", "cosine": "cosine", "normalized": "cosine"}[metric] + self.M = M + self.efConstruction = efConstruction + self.ef_query = None + + def _new_handle(self, dim): + return self.lib.hann_hnsw_new( + dim, self.M, self.efConstruction, self.metric.encode("ascii") + ) + + def _push_ef(self, ef): + if self.lib.hann_hnsw_set_ef(self.handle, ef) != 0: + raise RuntimeError("hann: failed to set ef to %d" % ef) + self._ef_current = ef + + def set_query_arguments(self, ef): + self._push_ef(ef) + self.ef_query = ef + + def query(self, v, n): + # Hann searches with the configured ef even when it is below k, and + # then completes the result with an exact scan. hnswlib instead + # searches with max(ef, k). Clamp the same way, so low-ef points + # measure the graph rather than the scan. + if self._ef_current < n: + self._push_ef(n) + return super().query(v, n) + + def __str__(self): + return "HannHNSW(M=%d, efConstruction=%d, efQuery=%d)" % ( + self.M, + self.efConstruction, + self.ef_query, + ) + + +class HannPQIVF(_HannBase): + def __init__(self, metric, coarseK, candidateClusters): + # The PQIVF index is Euclidean only. Euclidean ordering on unit + # vectors equals cosine ordering, so cosine data is normalized here + # and "normalized" data is already unit length. + if metric not in ("euclidean", "cosine", "normalized"): + raise NotImplementedError(f"HannPQIVF does not support metric {metric}") + super().__init__() + self.normalize = metric == "cosine" + self.coarseK = coarseK + self.candidateClusters = candidateClusters + + def _prepare(self, X): + if not self.normalize: + return X + X = np.asarray(X, dtype=np.float32) + norms = np.linalg.norm(X, axis=1, keepdims=True) + norms[norms == 0] = 1.0 + return X / norms + + def _new_handle(self, dim): + # Zeros keep the library defaults for numSubquantizers, pqK, and + # kMeansIters. + return self.lib.hann_pqivf_new(dim, self.coarseK, 0, 0, 0, self.candidateClusters) + + def fit(self, X): + super().fit(X) + if self.lib.hann_train(self.handle) != 0: + raise RuntimeError("hann: training failed") + + def __str__(self): + return "HannPQIVF(coarseK=%d, candidateClusters=%d)" % ( + self.coarseK, + self.candidateClusters, + ) + + +class HannRPT(_HannBase): + def __init__(self, metric, leafCapacity, probeMargin): + if metric not in ("euclidean", "cosine", "normalized"): + raise NotImplementedError(f"HannRPT does not support metric {metric}") + super().__init__() + self.metric = {"euclidean": "euclidean", "cosine": "cosine", "normalized": "cosine"}[metric] + self.leafCapacity = leafCapacity + self.probeMargin = probeMargin + + def _new_handle(self, dim): + # A zero keeps the library default for candidateProjections. + return self.lib.hann_rpt_new( + dim, self.leafCapacity, 0, self.probeMargin, self.metric.encode("ascii") + ) + + def __str__(self): + return "HannRPT(leafCapacity=%d, probeMargin=%g)" % ( + self.leafCapacity, + self.probeMargin, + ) diff --git a/vibe/results.py b/vibe/results.py index 58493da..7689a08 100644 --- a/vibe/results.py +++ b/vibe/results.py @@ -62,6 +62,10 @@ def build_result_filepath( d.append(definition.algorithm + suffix) index_parameters = re.sub(r"\W+", "_", json.dumps(definition.arguments, sort_keys=True)).strip("_") search_parameters = re.sub(r"\W+", "_", json.dumps(query_arguments, sort_keys=True)).strip("_") + if not search_parameters: + # An algorithm without query arguments still needs a non-empty + # HDF5 group name. + search_parameters = "default" d.append(index_parameters + ".hdf5") return os.path.join(*d), search_parameters