diff --git a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index dffb500..1e98b86 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -13,81 +13,24 @@ import java.util.concurrent.Executors; import java.util.function.Supplier; -public class AcceleratedHNSWParams { +public class AcceleratedHNSWParams extends ParamsBase { - public static enum Strategy { - /* - * This strategy allows for automatic selection of the underlying CAGRA build algorithm. - * With this strategy we use NN_DESCENT for dataset less than 5M vectors, else we use IVF_PQ. - * Indexing parameters, especially for IVF_PQ, are heuristically identified automatically. - * - * This is the default and the recommended strategy. - */ - HEURISTIC, - /* - * This is an option when the end-user would want to use custom parameter values. - * This strategy should only be used under expert guidance. - */ - CUSTOM - } - - /* - * TODO: Update boundaries for all parameters when a consensus is reached. - * Issue: https://github.com/rapidsai/cuvs-lucene/issues/99 - */ - public static final int MIN_WRITER_THREADS = 1; - public static final int MAX_WRITER_THREADS = 512; - public static final int MIN_INT_GRAPH_DEG = 2; - public static final int MAX_INT_GRAPH_DEG = 512; - public static final int MIN_GRAPH_DEG = 1; - public static final int MAX_GRAPH_DEG = 512; public static final int MIN_HNSW_LAYERS = 1; - public static final int MAX_HNSW_LAYERS = 3; - public static final int MIN_MAX_CONN = 1; - public static final int MAX_MAX_CONN = 512; - public static final int MIN_BEAM_WIDTH = 1; - public static final int MAX_BEAM_WIDTH = 512; + public static final int MAX_HNSW_LAYERS = 512; public static final int MIN_NUM_MERGE_WORKERS = 1; public static final int MAX_NUM_MERGE_WORKERS = 512; - public static final int MIN_NN_DESCENT_NUM_ITERATIONS = 1; - public static final int MAX_NN_DESCENT_NUM_ITERATIONS = 100; - public static final int DEFAULT_WRITER_THREADS = 1; - public static final int DEFAULT_INT_GRAPH_DEGREE = 128; - public static final int DEFAULT_GRAPH_DEGREE = 64; public static final int DEFAULT_HNSW_LAYERS = 1; - public static final int DEFAULT_MAX_CONN = 32; - public static final int DEFAULT_BEAM_WIDTH = 32; - public static final CagraGraphBuildAlgo DEFAULT_CAGRA_GRAPH_BUILD_ALGO = - CagraGraphBuildAlgo.NN_DESCENT; - public static final int DEFAULT_NUM_MERGE_WORKERS = 1; - public static final Strategy DEFAULT_STRATEGY = Strategy.HEURISTIC; - public static final CuvsDistanceType DEFAULT_CUVS_DISTANCE_TYPE = CuvsDistanceType.L2Expanded; - public static final int DEFAULT_NN_DESCENT_NUM_ITERATIONS = 20; - - public static final Supplier DEFAULT_IVF_PQ_PARAMS = - () -> { - return new CuVSIvfPqParams.Builder().build(); - }; + public static final int DEFAULT_NUM_MERGE_WORKERS = 16; public static final Supplier DEFAULT_MERGE_EXE_SRVC = () -> { return Executors.newFixedThreadPool(DEFAULT_NUM_MERGE_WORKERS); }; - private final int writerThreads; - private final int intermediateGraphDegree; - private final int graphdegree; private final int hnswLayers; - private final int maxConn; - private final int beamWidth; - private final CagraGraphBuildAlgo cagraGraphBuildAlgo; - private final CuVSIvfPqParams cuVSIvfPqParams; private final int numMergeWorkers; private final ExecutorService mergeExec; - private final Strategy strategy; - private final CuvsDistanceType cuvsDistanceType; - private final int nnDescentNumIterations; /** * Constructs an instance of {@link AcceleratedHNSWParams} with specific parameter values. @@ -121,47 +64,21 @@ private AcceleratedHNSWParams( Strategy strategy, CuvsDistanceType cuvsDistanceType, int nnDescentNumIterations) { - super(); - this.writerThreads = writerThreads; - this.intermediateGraphDegree = intermediateGraphDegree; - this.graphdegree = graphdegree; + + super( + writerThreads, + intermediateGraphDegree, + graphdegree, + maxConn, + beamWidth, + cagraGraphBuildAlgo, + cuVSIvfPqParams, + strategy, + cuvsDistanceType, + nnDescentNumIterations); this.hnswLayers = hnswLayers; - this.maxConn = maxConn; - this.beamWidth = beamWidth; - this.cagraGraphBuildAlgo = cagraGraphBuildAlgo; - this.cuVSIvfPqParams = cuVSIvfPqParams; this.numMergeWorkers = numMergeWorkers; this.mergeExec = mergeExec; - this.strategy = strategy; - this.cuvsDistanceType = cuvsDistanceType; - this.nnDescentNumIterations = nnDescentNumIterations; - } - - /** - * Get the cuVS writer threads parameter - * - * @return cuVS writer threads parameter - */ - public int getWriterThreads() { - return writerThreads; - } - - /** - * Get the intermediate graph degree - * - * @return the graph degree parameter - */ - public int getIntermediateGraphDegree() { - return intermediateGraphDegree; - } - - /** - * Get the graph degree - * - * @return the graph degree parameter - */ - public int getGraphdegree() { - return graphdegree; } /** @@ -173,42 +90,6 @@ public int getHnswLayers() { return hnswLayers; } - /** - * Get the max connection parameter - * - * @return the max connection parameter - */ - public int getMaxConn() { - return maxConn; - } - - /** - * Get the beam width parameter - * - * @return the beam width parameter - */ - public int getBeamWidth() { - return beamWidth; - } - - /** - * Get the CAGRA graph build algorithm - * - * @return the CAGRA graph build algorithm - */ - public CagraGraphBuildAlgo getCagraGraphBuildAlgo() { - return cagraGraphBuildAlgo; - } - - /** - * Get the instance of {@link CuVSIvfPqParams} - * - * @return the instance of {@link CuVSIvfPqParams} - */ - public CuVSIvfPqParams getCuVSIvfPqParams() { - return cuVSIvfPqParams; - } - /** * Get the number of merge workers set to be used in the fallback mechanism * @@ -227,64 +108,15 @@ public ExecutorService getMergeExec() { return mergeExec; } - /** - * Get the chosen strategy: - * - * When HEURISTIC [Default] is chosen, the CAGRA build algorithm and its indexing parameters are automatically chosen based on the size of the data set - * When CUSTOM is chosen, the build algorithm and its parameters (either defaults or overridden values with the use of With* methods) is used internally - * - * @return get the chosen {@link Strategy} - */ - public Strategy getStrategy() { - return strategy; - } - - /** - * Get the cuvs distance type - * - * @return the distance type - */ - public CuvsDistanceType getCuvsDistanceType() { - return cuvsDistanceType; - } - - /** - * get the number of Iterations to run if building with NN_DESCENT - * - * @return the number of iterations for NN_DESCENT - */ - public int getNNDescentNumIterations() { - return nnDescentNumIterations; - } - @Override public String toString() { - return "AcceleratedHNSWParams [writerThreads=" - + writerThreads - + ", intermediateGraphDegree=" - + intermediateGraphDegree - + ", graphdegree=" - + graphdegree - + ", hnswLayers=" + return super.toString() + + " AcceleratedHNSWParams [hnswLayers=" + hnswLayers - + ", maxConn=" - + maxConn - + ", beamWidth=" - + beamWidth - + ", cagraGraphBuildAlgo=" - + cagraGraphBuildAlgo - + ", cuVSIvfPqParams=" - + cuVSIvfPqParams + ", numMergeWorkers=" + numMergeWorkers + ", mergeExec=" + mergeExec - + ", strategy=" - + strategy - + ", cuvsDistanceType=" - + cuvsDistanceType - + ", nnDescentNumIterations=" - + nnDescentNumIterations + "]"; } @@ -480,31 +312,6 @@ public Builder withNNDescentNumIterations(int nnDescentNumIterations) { * @throws IllegalArgumentException */ private void validate() throws IllegalArgumentException { - if (writerThreads < MIN_WRITER_THREADS || writerThreads > MAX_WRITER_THREADS) { - throw new IllegalArgumentException( - "writerThreads not in valid range. Valid range: [" - + MIN_WRITER_THREADS - + ", " - + MAX_WRITER_THREADS - + "]"); - } - if (intermediateGraphDegree < MIN_INT_GRAPH_DEG - || intermediateGraphDegree > MAX_INT_GRAPH_DEG) { - throw new IllegalArgumentException( - "intermediateGraphDegree not in valid range. Valid range: [" - + MIN_INT_GRAPH_DEG - + ", " - + MAX_INT_GRAPH_DEG - + "]"); - } - if (graphdegree < MIN_GRAPH_DEG || graphdegree > MAX_GRAPH_DEG) { - throw new IllegalArgumentException( - "graphdegree not in valid range. Valid range: [" - + MIN_GRAPH_DEG - + ", " - + MAX_GRAPH_DEG - + "]"); - } if (hnswLayers < MIN_HNSW_LAYERS || hnswLayers > MAX_HNSW_LAYERS) { throw new IllegalArgumentException( "hnswLayers not in valid range. Valid range: [" @@ -513,25 +320,6 @@ private void validate() throws IllegalArgumentException { + MAX_HNSW_LAYERS + "]"); } - if (maxConn < MIN_MAX_CONN || maxConn > MAX_MAX_CONN) { - throw new IllegalArgumentException( - "maxConn not in valid range. Valid range: [" - + MIN_MAX_CONN - + ", " - + MAX_MAX_CONN - + "]"); - } - if (beamWidth < MIN_BEAM_WIDTH || beamWidth > MAX_BEAM_WIDTH) { - throw new IllegalArgumentException( - "beamWidth not in valid range. Valid range: [" - + MIN_BEAM_WIDTH - + ", " - + MAX_BEAM_WIDTH - + "]"); - } - if (Objects.isNull(cagraGraphBuildAlgo)) { - throw new IllegalArgumentException("cagraGraphBuildAlgo cannot be null."); - } if (numMergeWorkers < MIN_NUM_MERGE_WORKERS || numMergeWorkers > MAX_NUM_MERGE_WORKERS) { throw new IllegalArgumentException( "numMergeWorkers not in valid range. Valid range: [" @@ -540,21 +328,6 @@ private void validate() throws IllegalArgumentException { + MAX_NUM_MERGE_WORKERS + "]"); } - if (Objects.isNull(strategy)) { - throw new IllegalArgumentException("strategy cannot be null."); - } - if (Objects.isNull(cuvsDistanceType)) { - throw new IllegalArgumentException("cuvsDistanceType cannot be null."); - } - if (nnDescentNumIterations < MIN_NN_DESCENT_NUM_ITERATIONS - || nnDescentNumIterations > MAX_NN_DESCENT_NUM_ITERATIONS) { - throw new IllegalArgumentException( - "nnDescentNumIterations not in valid range. Valid range: [" - + MIN_NN_DESCENT_NUM_ITERATIONS - + ", " - + MAX_NN_DESCENT_NUM_ITERATIONS - + "]"); - } } /** @@ -563,9 +336,6 @@ private void validate() throws IllegalArgumentException { * @return instance of {@link AcceleratedHNSWParams} */ public AcceleratedHNSWParams build() { - if (Objects.isNull(cuVSIvfPqParams)) { - cuVSIvfPqParams = DEFAULT_IVF_PQ_PARAMS.get(); - } if (Objects.isNull(mergeExec)) { mergeExec = DEFAULT_MERGE_EXE_SRVC.get(); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 2a8ab02..da25546 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -90,8 +90,7 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( QuantizationType quantization) throws Throwable { - // Calculate M as cagraGraphDegree/2 - int M = graphDegree / 2; + int M = Math.ceilDiv(graphDegree, 2); // Store all layers data List layerNodes = new ArrayList<>(); @@ -320,7 +319,7 @@ public static void writeMeta( meta.writeVLong(vectorIndexLength); meta.writeVInt(field.getVectorDimension()); meta.writeInt(count); - meta.writeVInt(graphDegree / 2); // M = cagraGraphDegree/2 + meta.writeVInt(Math.ceilDiv(graphDegree, 2)); // write graph nodes on each level if (graph == null) { diff --git a/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java b/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java deleted file mode 100644 index 20fed04..0000000 --- a/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java +++ /dev/null @@ -1,225 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.nvidia.cuvs.lucene; - -import com.nvidia.cuvs.CagraIndexParams; -import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; -import com.nvidia.cuvs.CagraIndexParams.CodebookGen; -import com.nvidia.cuvs.CagraIndexParams.CudaDataType; -import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; -import com.nvidia.cuvs.CuVSIvfPqIndexParams; -import com.nvidia.cuvs.CuVSIvfPqParams; -import com.nvidia.cuvs.CuVSIvfPqSearchParams; - -/** - * A centralized approach to producing instances of {@link CagraIndexParams} based on the chosen strategy - */ -public class CagraIndexParamsFactory { - - private static final int ALGO_SWITCH_THRESHOLD = 5_000_000; - - /** - * Translation of the internal logic found here: - * https://github.com/rapidsai/cuvs/blob/main/cpp/include/cuvs/neighbors/ivf_pq.hpp#L3385-L3428 - * - * Ideally we should hook into the internal API but this is currently replicated to avoid complications - * in other parts of code base. - */ - private static CuVSIvfPqParams getCuVSIvfPqParams(long rows, long dimension) { - - int pqDim; - int pqBits; - - if (dimension <= 32) { - pqDim = 16; - pqBits = 8; - } else { - pqBits = 4; - if (dimension <= 64) { - pqDim = 32; - } else if (dimension <= 128) { - pqDim = 64; - } else if (dimension <= 192) { - pqDim = 96; - } else { - pqDim = (int) roundUpSafe(dimension / 2, 128); - } - } - - int nLists = (int) Math.max(1, rows / 2000); - final int kmeansNIters = 10; - final double kMinPointsPerCluster = 32; - double minKmeansTrainsetPoints = kMinPointsPerCluster * nLists; - final double maxKmeansTrainsetFraction = 1.0; - double minKmeansTrainsetFraction = - Math.min(maxKmeansTrainsetFraction, minKmeansTrainsetPoints / rows); - double kmeansTrainsetFraction = - Math.clamp( - 1.0 / Math.sqrt(rows * 1e-5), minKmeansTrainsetFraction, maxKmeansTrainsetFraction); - final CodebookGen codebookKind = CodebookGen.PER_SUBSPACE; - int nProbes = (int) Math.round(Math.sqrt(nLists) / 20 + 4); - final int refinementRate = 1; - - CuVSIvfPqIndexParams cuVSIvfPqIndexParams = - new CuVSIvfPqIndexParams.Builder() - .withCodebookKind(codebookKind) - .withKmeansNIters(kmeansNIters) - .withKmeansTrainsetFraction(kmeansTrainsetFraction) - .withNLists(nLists) - .withPqBits(pqBits) - .withPqDim(pqDim) - .withAddDataOnBuild(true) - .withConservativeMemoryAllocation(true) - .build(); - - CuVSIvfPqSearchParams cuVSIvfPqSearchParams = - new CuVSIvfPqSearchParams.Builder() - .withLutDtype(CudaDataType.CUDA_R_16F) - .withInternalDistanceDtype(CudaDataType.CUDA_R_16F) - .withNProbes(nProbes) - .build(); - - CuVSIvfPqParams cuVSIvfPqParams = - new CuVSIvfPqParams.Builder() - .withCuVSIvfPqIndexParams(cuVSIvfPqIndexParams) - .withCuVSIvfPqSearchParams(cuVSIvfPqSearchParams) - .withRefinementRate(refinementRate) - .build(); - - return cuVSIvfPqParams; - } - - /* - * Rough translation from raft's internal utility found here: - * https://github.com/rapidsai/raft/blob/main/cpp/include/raft/util/integer_utils.hpp#L47-L56 - */ - private static long roundUpSafe(long numberToRound, long modulus) { - long remainder = numberToRound % modulus; - if (remainder == 0) { - return numberToRound; - } - long roundedUp = numberToRound - remainder + modulus; - return roundedUp; - } - - private static CagraIndexParams getNNDescentParams( - int graphDegree, - int intGraphDegree, - int writerThreads, - long nnDescentNumIterations, - CuvsDistanceType cuvsDistanceType) { - return new CagraIndexParams.Builder() - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) - .withGraphDegree(graphDegree) - .withIntermediateGraphDegree(intGraphDegree) - .withNNDescentNumIterations(nnDescentNumIterations) - .withNumWriterThreads(writerThreads) - .withMetric(cuvsDistanceType) - .build(); - } - - private static CagraIndexParams getIVFPQParams( - int graphDegree, - int intGraphDegree, - int writerThreads, - long rows, - long dimension, - CuvsDistanceType cuvsDistanceType) { - return new CagraIndexParams.Builder() - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.IVF_PQ) - .withCuVSIvfPqParams(getCuVSIvfPqParams(rows, dimension)) - .withNumWriterThreads(writerThreads) - .withIntermediateGraphDegree(intGraphDegree) - .withGraphDegree(graphDegree) - .withMetric(cuvsDistanceType) - .build(); - } - - /** - * Creates an instance of {@link CagraIndexParams} based on the chosen strategy in the {@link GPUSearchParams}. - * - * @param gPUSearchParams an instance of {@link GPUSearchParams} containing input params incoming via the build and search on the GPU API. - * @param rows number of vectors in the data set - * @param dimension the dimension of the vectors in the data set - * @return an instance of {@link CagraIndexParams} - */ - public static CagraIndexParams create( - GPUSearchParams gPUSearchParams, long rows, long dimension) { - if (gPUSearchParams.getStrategy().equals(GPUSearchParams.Strategy.HEURISTIC)) { - if (rows < ALGO_SWITCH_THRESHOLD) { - return getNNDescentParams( - gPUSearchParams.getGraphdegree(), - gPUSearchParams.getIntermediateGraphDegree(), - gPUSearchParams.getWriterThreads(), - gPUSearchParams.getnNDescentNumIterations(), - gPUSearchParams.getCuvsDistanceType()); - } else { - return getIVFPQParams( - gPUSearchParams.getGraphdegree(), - gPUSearchParams.getIntermediateGraphDegree(), - gPUSearchParams.getWriterThreads(), - rows, - dimension, - gPUSearchParams.getCuvsDistanceType()); - } - } else { - return new CagraIndexParams.Builder() - .withNumWriterThreads(gPUSearchParams.getWriterThreads()) - .withIntermediateGraphDegree(gPUSearchParams.getIntermediateGraphDegree()) - .withGraphDegree(gPUSearchParams.getGraphdegree()) - .withCagraGraphBuildAlgo(gPUSearchParams.getCagraGraphBuildAlgo()) - .withCuVSIvfPqParams(gPUSearchParams.getCuVSIvfPqParams()) - .withNNDescentNumIterations(gPUSearchParams.getnNDescentNumIterations()) - .build(); - } - } - - /* - * Ideally there should be just one create method instead of two. - * We should do that when both the input parameter classes can be unified in the future. - */ - - /** - * Creates an instance of {@link CagraIndexParams} based on the chosen strategy in the {@link AcceleratedHNSWParams}. - * - * @param acceleratedHNSWParams an instance of {@link AcceleratedHNSWParams} containing input params incoming via the build and search on the GPU API. - * @param rows number of vectors in the data set - * @param dimension the dimension of the vectors in the data set - * @return an instance of {@link CagraIndexParams} - */ - public static CagraIndexParams create( - AcceleratedHNSWParams acceleratedHNSWParams, long rows, long dimension) { - if (acceleratedHNSWParams.getStrategy().equals(AcceleratedHNSWParams.Strategy.HEURISTIC)) { - if (rows - < ALGO_SWITCH_THRESHOLD) { // TODO: maybe consider making this threshold configurable from - // outside later. - return getNNDescentParams( - acceleratedHNSWParams.getGraphdegree(), - acceleratedHNSWParams.getIntermediateGraphDegree(), - acceleratedHNSWParams.getWriterThreads(), - acceleratedHNSWParams.getNNDescentNumIterations(), - acceleratedHNSWParams.getCuvsDistanceType()); - } else { - return getIVFPQParams( - acceleratedHNSWParams.getGraphdegree(), - acceleratedHNSWParams.getIntermediateGraphDegree(), - acceleratedHNSWParams.getWriterThreads(), - rows, - dimension, - acceleratedHNSWParams.getCuvsDistanceType()); - } - } else { - return new CagraIndexParams.Builder() - .withNumWriterThreads(acceleratedHNSWParams.getWriterThreads()) - .withIntermediateGraphDegree(acceleratedHNSWParams.getIntermediateGraphDegree()) - .withGraphDegree(acceleratedHNSWParams.getGraphdegree()) - .withCagraGraphBuildAlgo(acceleratedHNSWParams.getCagraGraphBuildAlgo()) - .withCuVSIvfPqParams(acceleratedHNSWParams.getCuVSIvfPqParams()) - .withNNDescentNumIterations(acceleratedHNSWParams.getNNDescentNumIterations()) - .build(); - } - } -} diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index e8fb304..73f5d10 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -11,6 +11,7 @@ import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.VERSION_CURRENT; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.Utils.getCagraIndexParams; import static com.nvidia.cuvs.lucene.Utils.info; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; @@ -242,7 +243,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro */ private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { CagraIndexParams params = - CagraIndexParamsFactory.create(gpuSearchParams, dataset.size(), dataset.columns()); + getCagraIndexParams(gpuSearchParams, dataset.size(), dataset.columns()); CagraIndex index = CagraIndex.newBuilder(getCuVSResourcesInstance()) .withDataset(dataset) diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.java b/src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.java index cc44526..0dbca23 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.java @@ -10,63 +10,12 @@ import com.nvidia.cuvs.CuVSIvfPqParams; import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; import java.util.Objects; -import java.util.function.Supplier; -public class GPUSearchParams { +public class GPUSearchParams extends ParamsBase { - public static enum Strategy { - /* - * This strategy allows for automatic selection of the underlining CAGRA build algorithm. - * With this strategy we use NN_DESCENT for data set less then 5M vectors else we use IVF_PQ. - * Indexing parameters, especially for IVF_PQ, are heuristically identified automatically. - * - * This is the default and the recommended strategy. - */ - HEURISTIC, - /* - * This is an option when the end-user would want to use custom parameter values. - * This strategy should only be used under expert guidance. - */ - CUSTOM - } - - /* - * TODO: Update boundaries for all parameters when a consensus is reached. - * Issue: https://github.com/rapidsai/cuvs-lucene/issues/99 - */ - public static final int MIN_WRITER_THREADS = 1; - public static final int MAX_WRITER_THREADS = 512; - public static final int MIN_INT_GRAPH_DEG = 2; - public static final int MAX_INT_GRAPH_DEG = 512; - public static final int MIN_GRAPH_DEG = 1; - public static final int MAX_GRAPH_DEG = 512; - public static final int MIN_NN_DESCENT_NUM_ITERATIONS = 1; - public static final int MAX_NN_DESCENT_NUM_ITERATIONS = 100; - - public static final int DEFAULT_INT_GRAPH_DEGREE = 128; - public static final int DEFAULT_GRAPH_DEGREE = 64; - public static final CagraGraphBuildAlgo DEFAULT_CAGRA_GRAPH_BUILD_ALGO = - CagraGraphBuildAlgo.NN_DESCENT; public static final IndexType DEFAULT_INDEX_TYPE = IndexType.CAGRA; - public static final int DEFAULT_WRITER_THREADS = 32; - public static final Strategy DEFAULT_STRATEGY = Strategy.HEURISTIC; - public static final CuvsDistanceType DEFAULT_CUVS_DISTANCE_TYPE = CuvsDistanceType.L2Expanded; - public static final int DEFAULT_NN_DESCENT_NUM_ITERATIONS = 20; - public static final Supplier DEFAULT_IVF_PQ_PARAMS = - () -> { - return new CuVSIvfPqParams.Builder().build(); - }; - - private final int writerThreads; - private final int intermediateGraphDegree; - private final int graphdegree; - private final CagraGraphBuildAlgo cagraGraphBuildAlgo; private final IndexType indexType; - private final CuVSIvfPqParams cuVSIvfPqParams; - private final Strategy strategy; - private final CuvsDistanceType cuvsDistanceType; - private final int nnDescentNumIterations; /** * Constructs an instance of {@link GPUSearchParams} with specific parameter values. @@ -91,53 +40,21 @@ private GPUSearchParams( CuVSIvfPqParams cuVSIvfPqParams, Strategy strategy, CuvsDistanceType cuvsDistanceType, - int nnDescentNumIterations) { - super(); - this.writerThreads = writerThreads; - this.intermediateGraphDegree = intermediateGraphDegree; - this.graphdegree = graphdegree; - this.cagraGraphBuildAlgo = cagraGraphBuildAlgo; + int nnDescentNumIterations, + int maxConn, + int beamWidth) { + super( + writerThreads, + intermediateGraphDegree, + graphdegree, + maxConn, + beamWidth, + cagraGraphBuildAlgo, + cuVSIvfPqParams, + strategy, + cuvsDistanceType, + nnDescentNumIterations); this.indexType = indexType; - this.cuVSIvfPqParams = cuVSIvfPqParams; - this.strategy = strategy; - this.cuvsDistanceType = cuvsDistanceType; - this.nnDescentNumIterations = nnDescentNumIterations; - } - - /** - * Get the cuVS writer threads parameter - * - * @return cuVS writer threads parameter - */ - public int getWriterThreads() { - return writerThreads; - } - - /** - * Get the intermediate graph degree - * - * @return the graph degree parameter - */ - public int getIntermediateGraphDegree() { - return intermediateGraphDegree; - } - - /** - * Get the graph degree - * - * @return the graph degree parameter - */ - public int getGraphdegree() { - return graphdegree; - } - - /** - * Get the CAGRA build algorithm parameter value - * - * @return the CAGRA build algorithm parameter value - */ - public CagraGraphBuildAlgo getCagraGraphBuildAlgo() { - return cagraGraphBuildAlgo; } /** @@ -149,67 +66,9 @@ public IndexType getIndexType() { return indexType; } - /** - * Get the instance of CuVSIvfPqParams - * - * @return an instance of CuVSIvfPqParams - */ - public CuVSIvfPqParams getCuVSIvfPqParams() { - return cuVSIvfPqParams; - } - - /** - * Get the chosen strategy: - * - * When HEURISTIC [Default] is chosen, the CAGRA build algorithm and its indexing parameters are automatically chosen based on the size of the data set - * When CUSTOM is chosen, the build algorithm and its parameters (either defaults or overridden values with the use of With* methods) is used internally - * - * - * @return get the chosen {@link Strategy} - */ - public Strategy getStrategy() { - return strategy; - } - - /** - * Get the cuvs distance type - * - * @return the distance type - */ - public CuvsDistanceType getCuvsDistanceType() { - return cuvsDistanceType; - } - - /** - * get the number of Iterations to run if building with NN_DESCENT - * - * @return the number of iterations for NN_DESCENT - */ - public int getnNDescentNumIterations() { - return nnDescentNumIterations; - } - @Override public String toString() { - return "GPUSearchParams [writerThreads=" - + writerThreads - + ", intermediateGraphDegree=" - + intermediateGraphDegree - + ", graphdegree=" - + graphdegree - + ", cagraGraphBuildAlgo=" - + cagraGraphBuildAlgo - + ", indexType=" - + indexType - + ", cuVSIvfPqParams=" - + cuVSIvfPqParams - + ", strategy=" - + strategy - + ", cuvsDistanceType=" - + cuvsDistanceType - + ", nnDescentNumIterations=" - + nnDescentNumIterations - + "]"; + return super.toString() + " GPUSearchParams [indexType=" + indexType + "]"; } /** @@ -226,6 +85,8 @@ public static class Builder { private Strategy strategy = DEFAULT_STRATEGY; private CuvsDistanceType cuvsDistanceType = DEFAULT_CUVS_DISTANCE_TYPE; private int nnDescentNumIterations = DEFAULT_NN_DESCENT_NUM_ITERATIONS; + private int maxConn = DEFAULT_MAX_CONN; + private int beamWidth = DEFAULT_BEAM_WIDTH; /** * Set the number of cuVS writer threads while building the index @@ -343,58 +204,41 @@ public Builder withNNDescentNumIterations(int nnDescentNumIterations) { return this; } + /** + * Set the max connections parameter while building HNSW index with fallback mechanism + * Valid range - Minimum: {@value MIN_MAX_CONN}, Maximum: {@value MAX_MAX_CONN} + * Default value - {@value DEFAULT_MAX_CONN} + * + * @param maxConn the max connections parameter + * @return instance of {@link Builder} + */ + public Builder withMaxConn(int maxConn) { + this.maxConn = maxConn; + return this; + } + + /** + * Set the beam width parameter while building HNSW index with fallback mechanism + * Valid range - Minimum: {@value MIN_BEAM_WIDTH}, Maximum: {@value MAX_BEAM_WIDTH} + * Default value - {@value DEFAULT_BEAM_WIDTH} + * + * @param beamWidth the beam width parameter + * @return instance of {@link Builder} + */ + public Builder withBeamWidth(int beamWidth) { + this.beamWidth = beamWidth; + return this; + } + /** * Validates the input parameters. * * @throws IllegalArgumentException */ private void validate() throws IllegalArgumentException { - if (writerThreads < MIN_WRITER_THREADS || writerThreads > MAX_WRITER_THREADS) { - throw new IllegalArgumentException( - "writerThreads not in valid range. Valid range: [" - + MIN_WRITER_THREADS - + ", " - + MAX_WRITER_THREADS - + "]"); - } - if (intermediateGraphDegree < MIN_INT_GRAPH_DEG - || intermediateGraphDegree > MAX_INT_GRAPH_DEG) { - throw new IllegalArgumentException( - "intermediateGraphDegree not in valid range. Valid range: [" - + MIN_INT_GRAPH_DEG - + ", " - + MAX_INT_GRAPH_DEG - + "]"); - } - if (graphdegree < MIN_GRAPH_DEG || graphdegree > MAX_GRAPH_DEG) { - throw new IllegalArgumentException( - "graphdegree not in valid range. Valid range: [" - + MIN_GRAPH_DEG - + ", " - + MAX_GRAPH_DEG - + "]"); - } - if (Objects.isNull(cagraGraphBuildAlgo)) { - throw new IllegalArgumentException("cagraGraphBuildAlgo cannot be null."); - } if (Objects.isNull(indexType)) { throw new IllegalArgumentException("indexType cannot be null."); } - if (Objects.isNull(strategy)) { - throw new IllegalArgumentException("strategy cannot be null."); - } - if (Objects.isNull(cuvsDistanceType)) { - throw new IllegalArgumentException("cuvsDistanceType cannot be null."); - } - if (nnDescentNumIterations < MIN_NN_DESCENT_NUM_ITERATIONS - || nnDescentNumIterations > MAX_NN_DESCENT_NUM_ITERATIONS) { - throw new IllegalArgumentException( - "nnDescentNumIterations not in valid range. Valid range: [" - + MIN_NN_DESCENT_NUM_ITERATIONS - + ", " - + MAX_NN_DESCENT_NUM_ITERATIONS - + "]"); - } } /** @@ -403,9 +247,6 @@ private void validate() throws IllegalArgumentException { * @return instance of {@link GPUSearchParams} */ public GPUSearchParams build() { - if (Objects.isNull(cuVSIvfPqParams)) { - cuVSIvfPqParams = DEFAULT_IVF_PQ_PARAMS.get(); - } validate(); return new GPUSearchParams( writerThreads, @@ -416,7 +257,9 @@ public GPUSearchParams build() { cuVSIvfPqParams, strategy, cuvsDistanceType, - nnDescentNumIterations); + nnDescentNumIterations, + maxConn, + beamWidth); } } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 35209ce..e318fd4 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -17,6 +17,7 @@ import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; import static com.nvidia.cuvs.lucene.Utils.createListFromMergedVectors; +import static com.nvidia.cuvs.lucene.Utils.getCagraIndexParams; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; @@ -160,7 +161,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); CagraIndexParams params = - CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); + getCagraIndexParams(acceleratedHNSWParams, dataset.size(), dataset.columns()); CagraIndex cagraIndex = CagraIndex.newBuilder(getCuVSResourcesInstance()) @@ -178,7 +179,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro adjacencyListMatrix, vectors, acceleratedHNSWParams.getHnswLayers(), - acceleratedHNSWParams.getGraphdegree(), + (int) params.getGraphDegree(), params, QuantizationType.NONE); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); @@ -193,7 +194,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro size, hnswGraph, graphLevelNodeOffsets, - acceleratedHNSWParams.getGraphdegree()); + (int) params.getGraphDegree()); cagraIndex.close(); } catch (Throwable t) { Utils.handleThrowable(t); diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index d8a98cc..90ef71c 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -17,6 +17,7 @@ import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_NAME; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.Utils.getCagraIndexParams; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; @@ -164,7 +165,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw } CagraIndexParams params = - CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); + getCagraIndexParams(acceleratedHNSWParams, dataset.size(), dataset.columns()); CagraIndex cagraIndex = CagraIndex.newBuilder(getCuVSResourcesInstance()) diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 21b4be3..eb3e96a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -17,6 +17,7 @@ import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_NAME; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.Utils.getCagraIndexParams; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; @@ -190,7 +191,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE } CagraIndexParams params = - CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); + getCagraIndexParams(acceleratedHNSWParams, dataset.size(), dataset.columns()); CagraIndex cagraIndex = CagraIndex.newBuilder(getCuVSResourcesInstance()) diff --git a/src/main/java/com/nvidia/cuvs/lucene/ParamsBase.java b/src/main/java/com/nvidia/cuvs/lucene/ParamsBase.java new file mode 100644 index 0000000..2bc910d --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/ParamsBase.java @@ -0,0 +1,297 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import com.nvidia.cuvs.CuVSIvfPqParams; +import java.util.Objects; +import java.util.function.Supplier; + +public class ParamsBase { + + public static enum Strategy { + /* + * This strategy allows for automatic selection of the underlying CAGRA build algorithm. + * With this strategy we use NN_DESCENT for dataset less than 5M vectors, else we use IVF_PQ. + * Indexing parameters, especially for IVF_PQ, are heuristically identified automatically. + * + * This is the default and the recommended strategy. + */ + HEURISTIC, + /* + * This is an option when the end-user would want to use custom parameter values. + * This strategy should only be used under expert guidance. + */ + CUSTOM + } + + /* + * TODO: Update boundaries for all parameters when a consensus is reached. + * Issue: https://github.com/rapidsai/cuvs-lucene/issues/99 + */ + public static final int MIN_WRITER_THREADS = 1; + public static final int MAX_WRITER_THREADS = 512; + public static final int MIN_INT_GRAPH_DEG = 2; + public static final int MAX_INT_GRAPH_DEG = 512; + public static final int MIN_GRAPH_DEG = 1; + public static final int MAX_GRAPH_DEG = 512; + public static final int MIN_MAX_CONN = 1; + public static final int MAX_MAX_CONN = 512; + public static final int MIN_BEAM_WIDTH = 1; + public static final int MAX_BEAM_WIDTH = 512; + public static final int MIN_NN_DESCENT_NUM_ITERATIONS = 1; + public static final int MAX_NN_DESCENT_NUM_ITERATIONS = 100; + + public static final int DEFAULT_WRITER_THREADS = 1; + public static final int DEFAULT_INT_GRAPH_DEGREE = 128; + public static final int DEFAULT_GRAPH_DEGREE = 64; + public static final int DEFAULT_MAX_CONN = 32; + public static final int DEFAULT_BEAM_WIDTH = 32; + public static final CagraGraphBuildAlgo DEFAULT_CAGRA_GRAPH_BUILD_ALGO = + CagraGraphBuildAlgo.NN_DESCENT; + public static final Strategy DEFAULT_STRATEGY = Strategy.HEURISTIC; + public static final CuvsDistanceType DEFAULT_CUVS_DISTANCE_TYPE = CuvsDistanceType.L2Expanded; + public static final int DEFAULT_NN_DESCENT_NUM_ITERATIONS = 20; + + public static final Supplier DEFAULT_IVF_PQ_PARAMS = + () -> { + return new CuVSIvfPqParams.Builder().build(); + }; + + private final int writerThreads; + private final int intermediateGraphDegree; + private final int graphdegree; + private final int maxConn; + private final int beamWidth; + private final CagraGraphBuildAlgo cagraGraphBuildAlgo; + private final CuVSIvfPqParams cuVSIvfPqParams; + private final Strategy strategy; + private final CuvsDistanceType cuvsDistanceType; + private final int nnDescentNumIterations; + + /** + * Constructs an instance of {@link AcceleratedHNSWParams} with specific parameter values. + * + * @param writerThreads Number of cuVS writer threads to use. + * @param intermediateGraphDegree The intermediate graph degree while building the CAGRA index. + * @param graphdegree The graph degree to use while building the CAGRA index. + * @param maxConn The max connection parameter used when building HNSW index with the fallback mechanism. + * @param beamWidth The beam width parameter used when building HNSW index with the fallback mechanism. + * @param cagraGraphBuildAlgo The CAGRA graph build algorithm to use [NN_DESCENT, IVF_PQ]. + * @param cuVSIvfPqParams An instance of CuVSIvfPqParams containing IVF_PQ specific parameters. + * @param strategy either HEURISTIC [Default] that automatically chooses build algorithm and its parameters based on data set size or CUSTOM that uses the parameters passed though this class. + * @param cuvsDistanceType the cuvsDistanceType. The default option is L2Expanded. + * @param nnDescentNumIterations the number of Iterations to run if building with NN_DESCENT. + */ + protected ParamsBase( + int writerThreads, + int intermediateGraphDegree, + int graphdegree, + int maxConn, + int beamWidth, + CagraGraphBuildAlgo cagraGraphBuildAlgo, + CuVSIvfPqParams cuVSIvfPqParams, + Strategy strategy, + CuvsDistanceType cuvsDistanceType, + int nnDescentNumIterations) { + this.writerThreads = writerThreads; + this.intermediateGraphDegree = intermediateGraphDegree; + this.graphdegree = graphdegree; + this.maxConn = maxConn; + this.beamWidth = beamWidth; + this.cagraGraphBuildAlgo = cagraGraphBuildAlgo; + this.cuVSIvfPqParams = cuVSIvfPqParams; + this.strategy = strategy; + this.cuvsDistanceType = cuvsDistanceType; + this.nnDescentNumIterations = nnDescentNumIterations; + validate(); + if (Objects.isNull(cuVSIvfPqParams)) { + cuVSIvfPqParams = DEFAULT_IVF_PQ_PARAMS.get(); + } + } + + /** + * Get the cuVS writer threads parameter + * + * @return cuVS writer threads parameter + */ + public int getWriterThreads() { + return writerThreads; + } + + /** + * Get the intermediate graph degree + * + * @return the graph degree parameter + */ + public int getIntermediateGraphDegree() { + return intermediateGraphDegree; + } + + /** + * Get the graph degree + * + * @return the graph degree parameter + */ + public int getGraphdegree() { + return graphdegree; + } + + /** + * Get the max connection parameter + * + * @return the max connection parameter + */ + public int getMaxConn() { + return maxConn; + } + + /** + * Get the beam width parameter + * + * @return the beam width parameter + */ + public int getBeamWidth() { + return beamWidth; + } + + /** + * Get the CAGRA graph build algorithm + * + * @return the CAGRA graph build algorithm + */ + public CagraGraphBuildAlgo getCagraGraphBuildAlgo() { + return cagraGraphBuildAlgo; + } + + /** + * Get the instance of {@link CuVSIvfPqParams} + *private + * @return the instance of {@link CuVSIvfPqParams} + */ + public CuVSIvfPqParams getCuVSIvfPqParams() { + return cuVSIvfPqParams; + } + + /** + * Get the chosen strategy: + * + * When HEURISTIC [Default] is chosen, the CAGRA build algorithm and its indexing parameters are automatically chosen based on the size of the data set + * When CUSTOM is chosen, the build algorithm and its parameters (either defaults or overridden values with the use of With* methods) is used internally + * + * @return get the chosen {@link Strategy} + */ + public Strategy getStrategy() { + return strategy; + } + + /** + * Get the cuvs distance type + * + * @return the distance type + */ + public CuvsDistanceType getCuvsDistanceType() { + return cuvsDistanceType; + } + + /** + * get the number of Iterations to run if building with NN_DESCENT + * + * @return the number of iterations for NN_DESCENT + */ + public int getNNDescentNumIterations() { + return nnDescentNumIterations; + } + + @Override + public String toString() { + return "ParamsBase [writerThreads=" + + writerThreads + + ", intermediateGraphDegree=" + + intermediateGraphDegree + + ", graphdegree=" + + graphdegree + + ", maxConn=" + + maxConn + + ", beamWidth=" + + beamWidth + + ", cagraGraphBuildAlgo=" + + cagraGraphBuildAlgo + + ", cuVSIvfPqParams=" + + cuVSIvfPqParams + + ", strategy=" + + strategy + + ", cuvsDistanceType=" + + cuvsDistanceType + + ", nnDescentNumIterations=" + + nnDescentNumIterations + + "]"; + } + + /** + * Validates the base input parameters. + * + * @throws IllegalArgumentException + */ + private void validate() throws IllegalArgumentException { + if (writerThreads < MIN_WRITER_THREADS || writerThreads > MAX_WRITER_THREADS) { + throw new IllegalArgumentException( + "writerThreads not in valid range. Valid range: [" + + MIN_WRITER_THREADS + + ", " + + MAX_WRITER_THREADS + + "]"); + } + if (intermediateGraphDegree < MIN_INT_GRAPH_DEG + || intermediateGraphDegree > MAX_INT_GRAPH_DEG) { + throw new IllegalArgumentException( + "intermediateGraphDegree not in valid range. Valid range: [" + + MIN_INT_GRAPH_DEG + + ", " + + MAX_INT_GRAPH_DEG + + "]"); + } + if (graphdegree < MIN_GRAPH_DEG || graphdegree > MAX_GRAPH_DEG) { + throw new IllegalArgumentException( + "graphdegree not in valid range. Valid range: [" + + MIN_GRAPH_DEG + + ", " + + MAX_GRAPH_DEG + + "]"); + } + if (maxConn < MIN_MAX_CONN || maxConn > MAX_MAX_CONN) { + throw new IllegalArgumentException( + "maxConn not in valid range. Valid range: [" + MIN_MAX_CONN + ", " + MAX_MAX_CONN + "]"); + } + if (beamWidth < MIN_BEAM_WIDTH || beamWidth > MAX_BEAM_WIDTH) { + throw new IllegalArgumentException( + "beamWidth not in valid range. Valid range: [" + + MIN_BEAM_WIDTH + + ", " + + MAX_BEAM_WIDTH + + "]"); + } + if (Objects.isNull(cagraGraphBuildAlgo)) { + throw new IllegalArgumentException("cagraGraphBuildAlgo cannot be null."); + } + if (Objects.isNull(strategy)) { + throw new IllegalArgumentException("strategy cannot be null."); + } + if (Objects.isNull(cuvsDistanceType)) { + throw new IllegalArgumentException("cuvsDistanceType cannot be null."); + } + if (nnDescentNumIterations < MIN_NN_DESCENT_NUM_ITERATIONS + || nnDescentNumIterations > MAX_NN_DESCENT_NUM_ITERATIONS) { + throw new IllegalArgumentException( + "nnDescentNumIterations not in valid range. Valid range: [" + + MIN_NN_DESCENT_NUM_ITERATIONS + + ", " + + MAX_NN_DESCENT_NUM_ITERATIONS + + "]"); + } + } +} diff --git a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java index 9e259e2..1eb3d49 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java @@ -6,6 +6,7 @@ package com.nvidia.cuvs.lucene; import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.spi.CuVSProvider; import java.util.logging.Level; import java.util.logging.Logger; @@ -19,9 +20,15 @@ public class ThreadLocalCuVSResourcesProvider { private static final Logger log = Logger.getLogger(ThreadLocalCuVSResourcesProvider.class.getName()); private static final ThreadLocal cuVSResources; + private static final CuVSProvider CUVS_PROVIDER; static { cuVSResources = ThreadLocal.withInitial(() -> cuVSResourcesOrNull()); + CUVS_PROVIDER = CuVSProvider.provider(); + } + + public static CuVSProvider getCuVSProvider() { + return CUVS_PROVIDER; } /** diff --git a/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/src/main/java/com/nvidia/cuvs/lucene/Utils.java index c55aa7c..a15c16a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -4,10 +4,14 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSProvider; import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CagraIndexParams.HnswHeuristicType; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.lucene.ParamsBase.Strategy; import java.io.IOException; import java.time.Duration; import java.util.ArrayList; @@ -205,4 +209,34 @@ static void info(InfoStream infoStream, String component, String msg) { infoStream.message(component, msg); } } + + /** + * Creates an instance of {@link CagraIndexParams} based on the chosen strategy in the {@link GPUSearchParams}. + * + * @param params an instance of {@link GPUSearchParams} containing input params incoming via the build and search on the GPU API. + * @param rows number of vectors in the data set + * @param dimension the dimension of the vectors in the data set + * @return an instance of {@link CagraIndexParams} + */ + public static CagraIndexParams getCagraIndexParams(ParamsBase params, long rows, long dimension) { + if (params.getStrategy().equals(Strategy.HEURISTIC)) { + return getCuVSProvider() + .cagraIndexParamsFromHnswParams( + rows, + dimension, + params.getMaxConn(), + params.getBeamWidth(), + HnswHeuristicType.SIMILAR_SEARCH_PERFORMANCE, + params.getCuvsDistanceType()); + } else { + return new CagraIndexParams.Builder() + .withNumWriterThreads(params.getWriterThreads()) + .withIntermediateGraphDegree(params.getIntermediateGraphDegree()) + .withGraphDegree(params.getGraphdegree()) + .withCagraGraphBuildAlgo(params.getCagraGraphBuildAlgo()) + .withCuVSIvfPqParams(params.getCuVSIvfPqParams()) + .withNNDescentNumIterations(params.getNNDescentNumIterations()) + .build(); + } + } } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.java b/src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.java index 483b689..f00723e 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.java @@ -49,7 +49,7 @@ public void testGPUSearchParamsDefaultValues() { assertEquals(DEFAULT_INDEX_TYPE, params.getIndexType()); assertEquals(DEFAULT_STRATEGY, params.getStrategy()); assertEquals(DEFAULT_CUVS_DISTANCE_TYPE, params.getCuvsDistanceType()); - assertEquals(DEFAULT_NN_DESCENT_NUM_ITERATIONS, params.getnNDescentNumIterations()); + assertEquals(DEFAULT_NN_DESCENT_NUM_ITERATIONS, params.getNNDescentNumIterations()); } @Test