From d0f53c0ca5f5c990a8c11fcb36e73e0a412de9de Mon Sep 17 00:00:00 2001 From: ldematte Date: Mon, 16 Jun 2025 17:28:21 +0200 Subject: [PATCH 01/11] Use Dataset to encapsulate on-heap arrays too --- .../main/java/com/nvidia/cuvs/Dataset.java | 15 +++++-- .../com/nvidia/cuvs/spi/CuVSProvider.java | 3 ++ .../nvidia/cuvs/spi/UnsupportedProvider.java | 5 +++ .../cuvs/internal/ArrayDatasetImpl.java | 42 +++++++++++++++++++ .../cuvs/internal/BruteForceIndexImpl.java | 21 ++++------ .../nvidia/cuvs/internal/CagraIndexImpl.java | 29 +++++-------- .../com/nvidia/cuvs/internal/DatasetImpl.java | 6 ++- .../cuvs/internal/MemorySegmentProvider.java | 8 ++++ .../com/nvidia/cuvs/spi/JDKProvider.java | 6 +++ 9 files changed, 99 insertions(+), 36 deletions(-) create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/ArrayDatasetImpl.java create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/MemorySegmentProvider.java diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java index 6fbbfee098..eac1d845d7 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java @@ -28,12 +28,21 @@ */ public interface Dataset extends AutoCloseable { + /** + * Creates a dataset from a on-heap array of vectors + * + * @since 25.08 + */ + static Dataset of(float[][] vectors) { + return CuVSProvider.provider().newArrayDataset(vectors); + } + /** * Add a single vector to the dataset. * * @param vector A float array of as many elements as the dimensions */ - public void addVector(float[] vector); + void addVector(float[] vector); /** * Create a new instance of a dataset @@ -51,12 +60,12 @@ static Dataset create(int size, int dimensions) { * * @return Size of the dataset */ - public int size(); + int size(); /** * Gets the dimensions of the vectors in this dataset * * @return Dimensions of the vectors in the dataset */ - public int dimensions(); + int dimensions(); } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java index a0d0f9a55c..3d4c1b923d 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java @@ -55,6 +55,9 @@ CuVSResources newCuVSResources(Path tempDirectory) /** Create a {@link Dataset} instance **/ Dataset newDataset(int size, int dimensions) throws UnsupportedOperationException; + /** Create a {@link Dataset} backed by a on-heap array **/ + Dataset newArrayDataset(float[][] vectors); + /** Creates a new BruteForceIndex Builder. */ BruteForceIndex.Builder newBruteForceIndexBuilder(CuVSResources cuVSResources) throws UnsupportedOperationException; diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index 5fe372dce0..7977d9253a 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -58,4 +58,9 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) throws Throwable { public Dataset newDataset(int size, int dimensions) throws UnsupportedOperationException { throw new UnsupportedOperationException(); } + + @Override + public Dataset newArrayDataset(float[][] vectors) { + throw new UnsupportedOperationException(); + } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/ArrayDatasetImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/ArrayDatasetImpl.java new file mode 100644 index 0000000000..e6c97d3d9c --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/ArrayDatasetImpl.java @@ -0,0 +1,42 @@ +package com.nvidia.cuvs.internal; + +import com.nvidia.cuvs.Dataset; +import com.nvidia.cuvs.internal.common.Util; + +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.util.Objects; + +public class ArrayDatasetImpl implements Dataset, MemorySegmentProvider { + private final float[][] vectors; + + public ArrayDatasetImpl(float[][] vectors) { + this.vectors = Objects.requireNonNull(vectors); + if (vectors.length == 0) { + throw new IllegalArgumentException("vectors should not be empty"); + } + } + + @Override + public void addVector(float[] vector) { + + } + + @Override + public int size() { + return vectors.length; + } + + @Override + public int dimensions() { + return vectors[0].length; + } + + @Override + public void close() { } + + @Override + public MemorySegment asMemorySegment(Arena arena) { + return Util.buildMemorySegment(arena, vectors); + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java index 141f5b0ab4..c54f65829d 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java @@ -76,7 +76,6 @@ */ public class BruteForceIndexImpl implements BruteForceIndex { - private final float[][] vectors; private final Dataset dataset; private final CuVSResourcesImpl resources; private final IndexReference bruteForceIndexReference; @@ -92,11 +91,10 @@ public class BruteForceIndexImpl implements BruteForceIndex { * @param bruteForceIndexParams an instance of {@link BruteForceIndexParams} * holding the index parameters */ - private BruteForceIndexImpl(float[][] vectors, Dataset dataset, CuVSResourcesImpl resources, + private BruteForceIndexImpl(Dataset dataset, CuVSResourcesImpl resources, BruteForceIndexParams bruteForceIndexParams) throws Throwable { - this.vectors = vectors; - this.dataset = dataset; + this.dataset = Objects.requireNonNull(dataset); this.resources = resources; this.bruteForceIndexParams = bruteForceIndexParams; this.bruteForceIndexReference = build(); @@ -110,7 +108,6 @@ private BruteForceIndexImpl(float[][] vectors, Dataset dataset, CuVSResourcesImp */ private BruteForceIndexImpl(InputStream inputStream, CuVSResourcesImpl resources) throws Throwable { this.bruteForceIndexParams = null; - this.vectors = null; this.dataset = null; this.resources = resources; this.bruteForceIndexReference = deserialize(inputStream); @@ -145,14 +142,13 @@ public void destroyIndex() throws Throwable { * @return an instance of {@link IndexReference} that holds the pointer to the * index */ - private IndexReference build() throws Throwable { + private IndexReference build() { try (var localArena = Arena.ofConfined()) { - long rows = dataset != null? dataset.size(): vectors.length; - long cols = dataset != null? dataset.dimensions(): (rows > 0 ? vectors[0].length : 0); + long rows = dataset.size(); + long cols = dataset.dimensions(); Arena arena = resources.getArena(); - MemorySegment datasetMemSegment = dataset != null? ((DatasetImpl) dataset).seg: - Util.buildMemorySegment(resources.getArena(), vectors); + MemorySegment datasetMemSegment = ((MemorySegmentProvider) dataset).asMemorySegment(arena); long cuvsResources = resources.getMemorySegment().get(cuvsResources_t, 0); MemorySegment stream = arena.allocate(cudaStream_t); @@ -383,7 +379,6 @@ public static BruteForceIndex.Builder newBuilder(CuVSResources cuvsResources) { */ public static class Builder implements BruteForceIndex.Builder { - private float[][] vectors; private Dataset dataset; private final CuVSResourcesImpl cuvsResources; private BruteForceIndexParams bruteForceIndexParams; @@ -432,7 +427,7 @@ public Builder from(InputStream inputStream) { */ @Override public Builder withDataset(float[][] vectors) { - this.vectors = vectors; + this.dataset = Dataset.of(vectors); return this; } @@ -458,7 +453,7 @@ public BruteForceIndexImpl build() throws Throwable { if (inputStream != null) { return new BruteForceIndexImpl(inputStream, cuvsResources); } else { - return new BruteForceIndexImpl(vectors, dataset, cuvsResources, bruteForceIndexParams); + return new BruteForceIndexImpl(dataset, cuvsResources, bruteForceIndexParams); } } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 33424e65c4..f1d11bc57f 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -97,7 +97,6 @@ * @since 25.02 */ public class CagraIndexImpl implements CagraIndex { - private final float[][] vectors; private final Dataset dataset; private final CuVSResourcesImpl resources; private final CagraIndexParams cagraIndexParameters; @@ -112,11 +111,9 @@ public class CagraIndexImpl implements CagraIndex { * @param dataset the dataset for indexing * @param resources an instance of {@link CuVSResources} */ - private CagraIndexImpl(CagraIndexParams indexParameters, float[][] vectors, - Dataset dataset, CuVSResourcesImpl resources) throws Throwable { + private CagraIndexImpl(CagraIndexParams indexParameters, Dataset dataset, CuVSResourcesImpl resources) { this.cagraIndexParameters = indexParameters; - this.vectors = vectors; - this.dataset = dataset; + this.dataset = Objects.requireNonNull(dataset); this.resources = resources; this.cagraIndexReference = build(); } @@ -129,7 +126,6 @@ private CagraIndexImpl(CagraIndexParams indexParameters, float[][] vectors, */ private CagraIndexImpl(InputStream inputStream, CuVSResourcesImpl resources) throws Throwable { this.cagraIndexParameters = null; - this.vectors = null; this.dataset = null; this.resources = resources; this.cagraIndexReference = deserialize(inputStream); @@ -143,7 +139,6 @@ private CagraIndexImpl(InputStream inputStream, CuVSResourcesImpl resources) thr * @param resources The resources instance */ private CagraIndexImpl(IndexReference indexReference, CuVSResourcesImpl resources) { - this.vectors = null; this.cagraIndexParameters = null; this.dataset = null; this.resources = resources; @@ -179,10 +174,10 @@ public void destroyIndex() throws Throwable { * @return an instance of {@link IndexReference} that holds the pointer to the * index */ - private IndexReference build() throws Throwable { + private IndexReference build() { try (var localArena = Arena.ofConfined()) { - long rows = dataset != null? dataset.size(): vectors.length; - long cols = dataset != null? dataset.dimensions(): (rows > 0 ? vectors[0].length : 0); + long rows = dataset.size(); + long cols = dataset.dimensions(); MemorySegment indexParamsMemorySegment = cagraIndexParameters != null ? segmentFromIndexParams(resources, cagraIndexParameters) @@ -191,10 +186,10 @@ private IndexReference build() throws Throwable { int numWriterThreads = cagraIndexParameters != null ? cagraIndexParameters.getNumWriterThreads() : 1; omp_set_num_threads(numWriterThreads); - MemorySegment dataSeg = dataset != null? ((DatasetImpl) dataset).seg: - Util.buildMemorySegment(resources.getArena(), vectors); - Arena arena = resources.getArena(); + assert dataset instanceof MemorySegmentProvider; + MemorySegment dataSeg = ((MemorySegmentProvider) dataset).asMemorySegment(arena); + long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); MemorySegment stream = arena.allocate(cudaStream_t); var returnValue = cuvsStreamGet(cuvsRes, stream); @@ -687,7 +682,6 @@ private static MemorySegment createMergeParamsSegment(CagraMergeParams mergePara */ public static class Builder implements CagraIndex.Builder { - private float[][] vectors; private Dataset dataset; private CagraIndexParams cagraIndexParams; private CuVSResourcesImpl cuvsResources; @@ -705,7 +699,7 @@ public Builder from(InputStream inputStream) { @Override public Builder withDataset(float[][] vectors) { - this.vectors = vectors; + this.dataset = Dataset.of(vectors); return this; } @@ -726,10 +720,7 @@ public CagraIndexImpl build() throws Throwable { if (inputStream != null) { return new CagraIndexImpl(inputStream, cuvsResources); } else { - if (vectors != null && dataset != null) { - throw new IllegalArgumentException("Please specify only one type of dataset (a float[] or a Dataset instance)"); - } - return new CagraIndexImpl(cagraIndexParams, vectors, dataset, cuvsResources); + return new CagraIndexImpl(cagraIndexParams, dataset, cuvsResources); } } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java index 13ac77e27e..faae132544 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java @@ -24,7 +24,7 @@ import com.nvidia.cuvs.Dataset; -public class DatasetImpl implements Dataset { +public class DatasetImpl implements Dataset, MemorySegmentProvider { private final Arena arena; protected final MemorySegment seg; private final int size; @@ -65,4 +65,8 @@ public int dimensions() { return dimensions; } + @Override + public MemorySegment asMemorySegment(Arena arena) { + return seg; + } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/MemorySegmentProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/MemorySegmentProvider.java new file mode 100644 index 0000000000..331a4d7128 --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/MemorySegmentProvider.java @@ -0,0 +1,8 @@ +package com.nvidia.cuvs.internal; + +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; + +interface MemorySegmentProvider { + MemorySegment asMemorySegment(Arena arena); +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index e7a2a73e27..56a3dec494 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -22,6 +22,7 @@ import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.HnswIndex; import com.nvidia.cuvs.CagraMergeParams; +import com.nvidia.cuvs.internal.ArrayDatasetImpl; import com.nvidia.cuvs.internal.BruteForceIndexImpl; import com.nvidia.cuvs.internal.CagraIndexImpl; import com.nvidia.cuvs.internal.CuVSResourcesImpl; @@ -81,4 +82,9 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams merge public Dataset newDataset(int size, int dimensions) throws UnsupportedOperationException { return new DatasetImpl(size, dimensions); } + + @Override + public Dataset newArrayDataset(float[][] vectors) { + return new ArrayDatasetImpl(vectors); + } } From 114285a546654c297a12172f76d3cf7c339b3bdc Mon Sep 17 00:00:00 2001 From: ldematte Date: Mon, 16 Jun 2025 17:45:48 +0200 Subject: [PATCH 02/11] Missing assert; renaming --- java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java | 2 +- .../java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java | 4 ++-- .../main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java | 3 +-- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java index eac1d845d7..aea3af5339 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java @@ -33,7 +33,7 @@ public interface Dataset extends AutoCloseable { * * @since 25.08 */ - static Dataset of(float[][] vectors) { + static Dataset ofArray(float[][] vectors) { return CuVSProvider.provider().newArrayDataset(vectors); } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java index c54f65829d..7da038a084 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java @@ -63,7 +63,6 @@ import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.SearchResults; -import com.nvidia.cuvs.internal.common.Util; import com.nvidia.cuvs.internal.panama.cuvsBruteForceIndex; import com.nvidia.cuvs.internal.panama.cuvsFilter; @@ -148,6 +147,7 @@ private IndexReference build() { long cols = dataset.dimensions(); Arena arena = resources.getArena(); + assert dataset instanceof MemorySegmentProvider; MemorySegment datasetMemSegment = ((MemorySegmentProvider) dataset).asMemorySegment(arena); long cuvsResources = resources.getMemorySegment().get(cuvsResources_t, 0); @@ -427,7 +427,7 @@ public Builder from(InputStream inputStream) { */ @Override public Builder withDataset(float[][] vectors) { - this.dataset = Dataset.of(vectors); + this.dataset = Dataset.ofArray(vectors); return this; } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index f1d11bc57f..c6b4c09fe9 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -72,7 +72,6 @@ import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.SearchResults; -import com.nvidia.cuvs.internal.common.Util; import com.nvidia.cuvs.internal.panama.cuvsCagraCompressionParams; import com.nvidia.cuvs.internal.panama.cuvsCagraIndex; import com.nvidia.cuvs.internal.panama.cuvsCagraIndexParams; @@ -699,7 +698,7 @@ public Builder from(InputStream inputStream) { @Override public Builder withDataset(float[][] vectors) { - this.dataset = Dataset.of(vectors); + this.dataset = Dataset.ofArray(vectors); return this; } From 02d58ab73e359f183664d42f694db5d5c3c49bde Mon Sep 17 00:00:00 2001 From: ldematte Date: Thu, 19 Jun 2025 12:36:31 +0200 Subject: [PATCH 03/11] Add benchmarks project --- java/.gitignore | 11 +- java/benchmarks/pom.xml | 154 ++++++++++++++++++ .../com/nvidia/cuvs/CagraIndexBenchmarks.java | 114 +++++++++++++ 3 files changed, 274 insertions(+), 5 deletions(-) create mode 100644 java/benchmarks/pom.xml create mode 100644 java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java diff --git a/java/.gitignore b/java/.gitignore index 2d3b14dd31..ab06477b43 100644 --- a/java/.gitignore +++ b/java/.gitignore @@ -1,9 +1,10 @@ +*.iml +target/ +jextract-22/ +openjdk-22-jextract* # cuvs-java -/cuvs-java/target/ /cuvs-java/bin/ /cuvs-java/src/main/java22/com/nvidia/cuvs/internal/panama/ /cuvs-java/*.cag -jextract-22/ -openjdk-22-jextract* -# examples -/examples/target/ + + diff --git a/java/benchmarks/pom.xml b/java/benchmarks/pom.xml new file mode 100644 index 0000000000..f5c77a14bc --- /dev/null +++ b/java/benchmarks/pom.xml @@ -0,0 +1,154 @@ + + + 4.0.0 + + com.nvidia.cuvs + benchmarks + 25.08.0 + jar + + cuvs-java-benchmarks + + + + org.openjdk.jmh + jmh-core + ${jmh.version} + + + org.openjdk.jmh + jmh-generator-annprocess + ${jmh.version} + provided + + + com.nvidia.cuvs + cuvs-java + 25.08.0 + jar + + + + + UTF-8 + 1.37 + 22 + benchmarks + + + + + + org.apache.maven.plugins + maven-compiler-plugin + 3.8.0 + + ${javac.target} + ${javac.target} + ${javac.target} + + -proc:full + + + + + org.apache.maven.plugins + maven-shade-plugin + 3.2.1 + + + package + + shade + + + ${uberjar.name} + + + org.openjdk.jmh.Main + + true + true + + + + + + + + *:* + + META-INF/*.SF + META-INF/*.DSA + META-INF/*.RSA + + + + + + + + + + + + maven-clean-plugin + 2.5 + + + maven-deploy-plugin + 2.8.1 + + + maven-install-plugin + 2.5.1 + + + maven-jar-plugin + 2.4 + + + maven-javadoc-plugin + 2.9.1 + + + maven-resources-plugin + 2.6 + + + maven-site-plugin + 3.3 + + + maven-source-plugin + 2.2.1 + + + maven-surefire-plugin + 2.17 + + + + + + diff --git a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java new file mode 100644 index 0000000000..15d47168b3 --- /dev/null +++ b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.nvidia.cuvs; + +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.infra.Blackhole; + +import java.lang.foreign.*; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Random; +import java.util.UUID; + +@Fork(value = 1, warmups = 0) +@State(Scope.Benchmark) +public class CagraIndexBenchmarks { + + @Param({ "1024" }) + private int dims; + + @Param({ "100" }) + private int size; + + private float[][] arrayDataset; + + private static final Random random = new Random(); + + private static float[][] createSampleData(int size, int dimensions) { + var array = new float[size][dimensions]; + for (int i = 0; i < size; ++i) { + for (int j = 0; j < dimensions; ++j) { + array[i][j] = random.nextFloat(); + } + } + return array; + } + + @Setup + public void initialize() { + arrayDataset = createSampleData(size, dims); + } + + @Benchmark + public void testIndexingAndSerializeToFile() throws Throwable { + try (CuVSResources resources = CuVSResources.create()) { + // Configure index parameters + CagraIndexParams indexParams = new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(1) + .withIntermediateGraphDegree(2) + .withNumWriterThreads(32) + .withMetric(CagraIndexParams.CuvsDistanceType.L2Expanded) + .build(); + + // Create the index with the dataset + CagraIndex index = CagraIndex.newBuilder(resources) + .withDataset(arrayDataset) + .withIndexParams(indexParams) + .build(); + + // Saving the index on to the disk. + var indexFilePath = Path.of(UUID.randomUUID() + ".cag"); + try (var outputStream = Files.newOutputStream(indexFilePath)) { + index.serialize(outputStream); + } + + // Cleanup + Files.deleteIfExists(indexFilePath); + index.destroyIndex(); + } + } + + @Benchmark + public void testIndexingFromHeap(Blackhole blackhole) throws Throwable { + try (CuVSResources resources = CuVSResources.create()) { + + // Configure index parameters + CagraIndexParams indexParams = new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(1) + .withIntermediateGraphDegree(2) + .withNumWriterThreads(32) + .withMetric(CagraIndexParams.CuvsDistanceType.L2Expanded) + .build(); + + // Create the index with the dataset + CagraIndex index = CagraIndex.newBuilder(resources) + .withDataset(arrayDataset) + .withIndexParams(indexParams) + .build(); + blackhole.consume(index); + } + } +} From 84ec3d95de801be78a6cf1ced9fb8eec8aea5bd5 Mon Sep 17 00:00:00 2001 From: ldematte Date: Thu, 19 Jun 2025 12:49:55 +0200 Subject: [PATCH 04/11] Add benchmarks README --- java/benchmarks/README.md | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 java/benchmarks/README.md diff --git a/java/benchmarks/README.md b/java/benchmarks/README.md new file mode 100644 index 0000000000..447c578448 --- /dev/null +++ b/java/benchmarks/README.md @@ -0,0 +1,27 @@ +# CuVS Java API benchmarks + +This maven project contains JMH benchmarks for the CAGRA Java API. + +## Prerequisites +- [CuVS libraries](https://docs.rapids.ai/api/cuvs/stable/build/#build-from-source) +- Build the CuVS-Java API (`./build.sh` from the parent directory) + +## Run benchmarks + +Build: +```shell +mvn clean verify +``` +Run: +```shell +export RAFT_DEBUG_LOG_FILE=/dev/null +java -jar target/benchmarks.jar +``` +The environment variable is needed to silence RAFT logging; RAFT emits some logs at INFO level when +building indices and queries, and writing them to stdout (the default) influences benchmark results. + +It is possible to change the dataset size and the vectors dimension via 2 parameters: +```shell +java -jar target/benchmarks.jar -p size=4 -p dims=4 +``` +Use `java -jar target/benchmarks.jar -h` for details on the options to fine-tune your benchmark runs. \ No newline at end of file From 01ca37e9103dccf1e339a9e8894b783e946929fd Mon Sep 17 00:00:00 2001 From: ldematte Date: Thu, 19 Jun 2025 13:02:04 +0200 Subject: [PATCH 05/11] Add MemorySegment based Dataset (0-copy) + benchmark --- .../com/nvidia/cuvs/CagraIndexBenchmarks.java | 47 +++++++++++++++++++ .../main/java/com/nvidia/cuvs/Dataset.java | 4 ++ .../com/nvidia/cuvs/spi/CuVSProvider.java | 2 + .../nvidia/cuvs/spi/UnsupportedProvider.java | 5 ++ .../com/nvidia/cuvs/internal/DatasetImpl.java | 7 +++ .../com/nvidia/cuvs/spi/JDKProvider.java | 7 +++ 6 files changed, 72 insertions(+) diff --git a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java index 15d47168b3..ae4fdc7276 100644 --- a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java +++ b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java @@ -43,6 +43,10 @@ public class CagraIndexBenchmarks { private float[][] arrayDataset; + private Arena arena; + + private MemorySegment memorySegmentDataset; + private static final Random random = new Random(); private static float[][] createSampleData(int size, int dimensions) { @@ -55,9 +59,31 @@ private static float[][] createSampleData(int size, int dimensions) { return array; } + private static MemorySegment createSampleDataSegment(Arena arena, float[][] array, int size, int dimensions) { + final ValueLayout.OfFloat C_FLOAT = (ValueLayout.OfFloat) Linker.nativeLinker().canonicalLayouts().get("float"); + + MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout((long)size * dimensions, C_FLOAT); + + var segment = arena.allocate(dataMemoryLayout); + for (int i = 0; i < size; ++i) { + var vector = array[i]; + MemorySegment.copy(vector, 0, segment, C_FLOAT, (i * dimensions * C_FLOAT.byteSize()), dimensions); + } + return segment; + } + @Setup public void initialize() { + arena = Arena.ofShared(); arrayDataset = createSampleData(size, dims); + memorySegmentDataset = createSampleDataSegment(arena, arrayDataset, size, dims); + } + + @TearDown + public void cleanUp() { + if (arena != null) { + arena.close(); + } } @Benchmark @@ -111,4 +137,25 @@ public void testIndexingFromHeap(Blackhole blackhole) throws Throwable { blackhole.consume(index); } } + + @Benchmark + public void testIndexingFromMemorySegment(Blackhole blackhole) throws Throwable { + try (CuVSResources resources = CuVSResources.create()) { + // Configure index parameters + CagraIndexParams indexParams = new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(1) + .withIntermediateGraphDegree(2) + .withNumWriterThreads(32) + .withMetric(CagraIndexParams.CuvsDistanceType.L2Expanded) + .build(); + + // Create the index with the dataset + CagraIndex index = CagraIndex.newBuilder(resources) + .withDataset(Dataset.ofMemorySegment(memorySegmentDataset, size, dims)) + .withIndexParams(indexParams) + .build(); + blackhole.consume(index); + } + } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java index aea3af5339..805917f894 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java @@ -37,6 +37,10 @@ static Dataset ofArray(float[][] vectors) { return CuVSProvider.provider().newArrayDataset(vectors); } + static Dataset ofMemorySegment(Object memorySegment, int size, int dimensions) { + return CuVSProvider.provider().newMemoryDataset(memorySegment, size, dimensions); + } + /** * Add a single vector to the dataset. * diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java index 3d4c1b923d..7695e578f5 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java @@ -55,6 +55,8 @@ CuVSResources newCuVSResources(Path tempDirectory) /** Create a {@link Dataset} instance **/ Dataset newDataset(int size, int dimensions) throws UnsupportedOperationException; + Dataset newMemoryDataset(Object memorySegment, int size, int dimensions); + /** Create a {@link Dataset} backed by a on-heap array **/ Dataset newArrayDataset(float[][] vectors); diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index 7977d9253a..21c1c4096b 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -59,6 +59,11 @@ public Dataset newDataset(int size, int dimensions) throws UnsupportedOperationE throw new UnsupportedOperationException(); } + @Override + public Dataset newMemoryDataset(Object memorySegment, int size, int dimensions) { + throw new UnsupportedOperationException(); + } + @Override public Dataset newArrayDataset(float[][] vectors) { throw new UnsupportedOperationException(); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java index faae132544..69486a04d9 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java @@ -41,6 +41,13 @@ public DatasetImpl(int size, int dimensions) { seg = arena.allocate(dataMemoryLayout); } + public DatasetImpl(MemorySegment memorySegment, int size, int dimensions) { + this.arena = null; + this.seg = memorySegment; + this.size = size; + this.dimensions = dimensions; + } + @Override public void addVector(float[] vector) { if (current >= size) diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 56a3dec494..c9b10cba8b 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -29,6 +29,7 @@ import com.nvidia.cuvs.internal.DatasetImpl; import com.nvidia.cuvs.internal.HnswIndexImpl; +import java.lang.foreign.MemorySegment; import java.nio.file.Files; import java.nio.file.Path; import java.util.Objects; @@ -83,6 +84,12 @@ public Dataset newDataset(int size, int dimensions) throws UnsupportedOperationE return new DatasetImpl(size, dimensions); } + @Override + public Dataset newMemoryDataset(Object memorySegment, int size, int dimensions) { + return new DatasetImpl((MemorySegment) memorySegment, size, dimensions); + } + + @Override public Dataset newArrayDataset(float[][] vectors) { return new ArrayDatasetImpl(vectors); From 86d95752ce57fcedf340ce47eb5ca1dab85220e2 Mon Sep 17 00:00:00 2001 From: ldematte Date: Thu, 19 Jun 2025 14:27:47 +0200 Subject: [PATCH 06/11] Alternative: move MemorySegment creation "up", at params build time --- .../cuvs/internal/ArrayDatasetImpl.java | 42 ------------------- .../cuvs/internal/BruteForceIndexImpl.java | 2 +- .../nvidia/cuvs/internal/CagraIndexImpl.java | 2 +- .../com/nvidia/cuvs/internal/DatasetImpl.java | 12 +++--- .../cuvs/internal/MemorySegmentProvider.java | 3 +- .../com/nvidia/cuvs/spi/JDKProvider.java | 17 ++++++-- 6 files changed, 23 insertions(+), 55 deletions(-) delete mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/ArrayDatasetImpl.java diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/ArrayDatasetImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/ArrayDatasetImpl.java deleted file mode 100644 index e6c97d3d9c..0000000000 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/ArrayDatasetImpl.java +++ /dev/null @@ -1,42 +0,0 @@ -package com.nvidia.cuvs.internal; - -import com.nvidia.cuvs.Dataset; -import com.nvidia.cuvs.internal.common.Util; - -import java.lang.foreign.Arena; -import java.lang.foreign.MemorySegment; -import java.util.Objects; - -public class ArrayDatasetImpl implements Dataset, MemorySegmentProvider { - private final float[][] vectors; - - public ArrayDatasetImpl(float[][] vectors) { - this.vectors = Objects.requireNonNull(vectors); - if (vectors.length == 0) { - throw new IllegalArgumentException("vectors should not be empty"); - } - } - - @Override - public void addVector(float[] vector) { - - } - - @Override - public int size() { - return vectors.length; - } - - @Override - public int dimensions() { - return vectors[0].length; - } - - @Override - public void close() { } - - @Override - public MemorySegment asMemorySegment(Arena arena) { - return Util.buildMemorySegment(arena, vectors); - } -} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java index 7da038a084..f2e40b431d 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java @@ -148,7 +148,7 @@ private IndexReference build() { Arena arena = resources.getArena(); assert dataset instanceof MemorySegmentProvider; - MemorySegment datasetMemSegment = ((MemorySegmentProvider) dataset).asMemorySegment(arena); + MemorySegment datasetMemSegment = ((MemorySegmentProvider) dataset).asMemorySegment(); long cuvsResources = resources.getMemorySegment().get(cuvsResources_t, 0); MemorySegment stream = arena.allocate(cudaStream_t); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index c6b4c09fe9..31ed6c906a 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -187,7 +187,7 @@ private IndexReference build() { Arena arena = resources.getArena(); assert dataset instanceof MemorySegmentProvider; - MemorySegment dataSeg = ((MemorySegmentProvider) dataset).asMemorySegment(arena); + MemorySegment dataSeg = ((MemorySegmentProvider) dataset).asMemorySegment(); long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); MemorySegment stream = arena.allocate(cudaStream_t); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java index 69486a04d9..b39987c28b 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java @@ -35,14 +35,14 @@ public DatasetImpl(int size, int dimensions) { this.size = size; this.dimensions = dimensions; - MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout(size * dimensions, C_FLOAT); + MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout((long)size * dimensions, C_FLOAT); this.arena = Arena.ofShared(); seg = arena.allocate(dataMemoryLayout); } - public DatasetImpl(MemorySegment memorySegment, int size, int dimensions) { - this.arena = null; + public DatasetImpl(Arena arena, MemorySegment memorySegment, int size, int dimensions) { + this.arena = arena; this.seg = memorySegment; this.size = size; this.dimensions = dimensions; @@ -52,12 +52,12 @@ public DatasetImpl(MemorySegment memorySegment, int size, int dimensions) { public void addVector(float[] vector) { if (current >= size) throw new ArrayIndexOutOfBoundsException(); - MemorySegment.copy(vector, 0, seg, C_FLOAT, ((current++) * dimensions * C_FLOAT.byteSize()), (int) dimensions); + MemorySegment.copy(vector, 0, seg, C_FLOAT, ((current++) * dimensions * C_FLOAT.byteSize()), dimensions); } @Override public void close() { - if (!arena.scope().isAlive()) { + if (arena != null && !arena.scope().isAlive()) { arena.close(); } } @@ -73,7 +73,7 @@ public int dimensions() { } @Override - public MemorySegment asMemorySegment(Arena arena) { + public MemorySegment asMemorySegment() { return seg; } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/MemorySegmentProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/MemorySegmentProvider.java index 331a4d7128..7ea460a1a8 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/MemorySegmentProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/MemorySegmentProvider.java @@ -1,8 +1,7 @@ package com.nvidia.cuvs.internal; -import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; interface MemorySegmentProvider { - MemorySegment asMemorySegment(Arena arena); + MemorySegment asMemorySegment(); } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index c9b10cba8b..67da34936e 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -22,13 +22,14 @@ import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.HnswIndex; import com.nvidia.cuvs.CagraMergeParams; -import com.nvidia.cuvs.internal.ArrayDatasetImpl; import com.nvidia.cuvs.internal.BruteForceIndexImpl; import com.nvidia.cuvs.internal.CagraIndexImpl; import com.nvidia.cuvs.internal.CuVSResourcesImpl; import com.nvidia.cuvs.internal.DatasetImpl; import com.nvidia.cuvs.internal.HnswIndexImpl; +import com.nvidia.cuvs.internal.common.Util; +import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; import java.nio.file.Files; import java.nio.file.Path; @@ -86,12 +87,22 @@ public Dataset newDataset(int size, int dimensions) throws UnsupportedOperationE @Override public Dataset newMemoryDataset(Object memorySegment, int size, int dimensions) { - return new DatasetImpl((MemorySegment) memorySegment, size, dimensions); + return new DatasetImpl(null, (MemorySegment) memorySegment, size, dimensions); } @Override public Dataset newArrayDataset(float[][] vectors) { - return new ArrayDatasetImpl(vectors); + Objects.requireNonNull(vectors); + if (vectors.length == 0) { + throw new IllegalArgumentException("vectors should not be empty"); + } + int size = vectors.length; + int dimensions = vectors[0].length; + + // TODO: should we use the Arena from CuVsResources here instead? + Arena arena = Arena.ofShared(); + var memorySegment = Util.buildMemorySegment(arena, vectors); + return new DatasetImpl(arena, memorySegment, size, dimensions); } } From 7f91a23acb4001f5ee751fd14573c98b0c10af7f Mon Sep 17 00:00:00 2001 From: ldematte Date: Thu, 19 Jun 2025 17:38:12 +0200 Subject: [PATCH 07/11] Different MemorySegment lifetime management (reduced scope) + specific Dataset benchmarks --- java/README.md | 8 +++- .../com/nvidia/cuvs/CagraIndexBenchmarks.java | 24 ++++++++++- .../main/java/com/nvidia/cuvs/CagraIndex.java | 7 ---- .../cuvs/internal/BruteForceIndexImpl.java | 28 ++++++------- .../nvidia/cuvs/internal/CagraIndexImpl.java | 42 ++++++------------- .../com/nvidia/cuvs/internal/DatasetImpl.java | 7 ++-- .../cuvs/internal/MemorySegmentProvider.java | 7 ---- .../com/nvidia/cuvs/spi/JDKProvider.java | 1 - .../nvidia/cuvs/CagraBuildAndSearchIT.java | 1 + 9 files changed, 59 insertions(+), 66 deletions(-) delete mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/MemorySegmentProvider.java diff --git a/java/README.md b/java/README.md index b568a1ffa0..b913ebe709 100644 --- a/java/README.md +++ b/java/README.md @@ -25,10 +25,14 @@ do `./build.sh java` in the top level directory or just do `./build.sh` in this Run `./build.sh --run-java-tests` from this directory. -To run a single test: +To run a single test suite: ```sh cd cuvs-java/ -mvn verify -Dintegration-test=com.nvidia.cuvs.CagraBuildAndSearchIT +mvn clean integration-test -Dit.test=com.nvidia.cuvs.CagraBuildAndSearchIT +``` +or, for a single test: +```sh +mvn clean integration-test -Dit.test=com.nvidia.cuvs.CagraBuildAndSearchIT#testMergeStrategies ``` Be sure to set (manually, if needed) your `LD_LIBRARY_PATH` to include the directory with the appropriate (matching) version of `libcuvs.so`. diff --git a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java index ae4fdc7276..8a0a97936f 100644 --- a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java +++ b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java @@ -17,12 +17,15 @@ package com.nvidia.cuvs; import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; import org.openjdk.jmh.annotations.Param; import org.openjdk.jmh.annotations.Scope; import org.openjdk.jmh.annotations.Setup; -import org.openjdk.jmh.annotations.TearDown; import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; import org.openjdk.jmh.infra.Blackhole; import java.lang.foreign.*; @@ -30,6 +33,7 @@ import java.nio.file.Path; import java.util.Random; import java.util.UUID; +import java.util.concurrent.TimeUnit; @Fork(value = 1, warmups = 0) @State(Scope.Benchmark) @@ -158,4 +162,22 @@ public void testIndexingFromMemorySegment(Blackhole blackhole) throws Throwable blackhole.consume(index); } } + + @Benchmark + @OutputTimeUnit(TimeUnit.NANOSECONDS) + @BenchmarkMode(Mode.AverageTime) + public void testDatasetFromHeap(Blackhole blackhole) throws Throwable { + try (var dataset = Dataset.ofArray(arrayDataset)) { + blackhole.consume(dataset); + } + } + + @Benchmark + @OutputTimeUnit(TimeUnit.NANOSECONDS) + @BenchmarkMode(Mode.AverageTime) + public void testDatasetFromMemorySegment(Blackhole blackhole) throws Throwable { + try (var dataset = Dataset.ofMemorySegment(memorySegmentDataset, size, dims)) { + blackhole.consume(dataset); + } + } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index a9a94df545..aae46dcb7a 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java @@ -145,13 +145,6 @@ default void serializeToHNSW(OutputStream outputStream, Path tempFile) throws Th */ void serializeToHNSW(OutputStream outputStream, Path tempFile, int bufferLength) throws Throwable; - /** - * Gets an instance of {@link CagraIndexParams} - * - * @return an instance of {@link CagraIndexParams} - */ - CagraIndexParams getCagraIndexParameters(); - /** * Gets an instance of {@link CuVSResources} * diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java index f2e40b431d..494f8dcee9 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java @@ -75,7 +75,6 @@ */ public class BruteForceIndexImpl implements BruteForceIndex { - private final Dataset dataset; private final CuVSResourcesImpl resources; private final IndexReference bruteForceIndexReference; private final BruteForceIndexParams bruteForceIndexParams; @@ -91,12 +90,14 @@ public class BruteForceIndexImpl implements BruteForceIndex { * holding the index parameters */ private BruteForceIndexImpl(Dataset dataset, CuVSResourcesImpl resources, - BruteForceIndexParams bruteForceIndexParams) - throws Throwable { - this.dataset = Objects.requireNonNull(dataset); - this.resources = resources; - this.bruteForceIndexParams = bruteForceIndexParams; - this.bruteForceIndexReference = build(); + BruteForceIndexParams bruteForceIndexParams) throws Exception { + Objects.requireNonNull(dataset); + try (dataset) { + this.resources = resources; + this.bruteForceIndexParams = bruteForceIndexParams; + assert dataset instanceof DatasetImpl; + this.bruteForceIndexReference = build((DatasetImpl) dataset); + } } /** @@ -107,7 +108,6 @@ private BruteForceIndexImpl(Dataset dataset, CuVSResourcesImpl resources, */ private BruteForceIndexImpl(InputStream inputStream, CuVSResourcesImpl resources) throws Throwable { this.bruteForceIndexParams = null; - this.dataset = null; this.resources = resources; this.bruteForceIndexReference = deserialize(inputStream); } @@ -131,7 +131,6 @@ public void destroyIndex() throws Throwable { } finally { destroyed = true; } - if (dataset != null) dataset.close(); } /** @@ -141,14 +140,13 @@ public void destroyIndex() throws Throwable { * @return an instance of {@link IndexReference} that holds the pointer to the * index */ - private IndexReference build() { + private IndexReference build(DatasetImpl dataset) { try (var localArena = Arena.ofConfined()) { long rows = dataset.size(); long cols = dataset.dimensions(); Arena arena = resources.getArena(); - assert dataset instanceof MemorySegmentProvider; - MemorySegment datasetMemSegment = ((MemorySegmentProvider) dataset).asMemorySegment(); + MemorySegment datasetMemSegment = dataset.asMemorySegment(); long cuvsResources = resources.getMemorySegment().get(cuvsResources_t, 0); MemorySegment stream = arena.allocate(cudaStream_t); @@ -168,7 +166,7 @@ private IndexReference build() { cudaMemcpy(datasetMemorySegmentP, datasetMemSegment, datasetBytes, INFER_DIRECTION); - long datasetShape[] = { rows, cols }; + long[] datasetShape = { rows, cols }; MemorySegment datasetTensor = prepareTensor(arena, datasetMemorySegmentP, datasetShape, 2, 32, 2, 2, 1); MemorySegment index = arena.allocate(cuvsBruteForceIndex_t); @@ -219,9 +217,9 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable { BitSet[] prefilters = cuvsQuery.getPrefilters(); if (prefilters != null && prefilters.length > 0) { BitSet concatenatedFilters = concatenate(prefilters, cuvsQuery.getNumDocs()); - long filters[] = concatenatedFilters.toLongArray(); + long[] filters = concatenatedFilters.toLongArray(); prefilterDataMemorySegment = buildMemorySegment(arena, filters); - prefilterDataLength = cuvsQuery.getNumDocs() * prefilters.length; + prefilterDataLength = (long)cuvsQuery.getNumDocs() * prefilters.length; } MemorySegment querySeg = buildMemorySegment(arena, cuvsQuery.getQueryVectors()); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 31ed6c906a..8563a998e4 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -96,9 +96,7 @@ * @since 25.02 */ public class CagraIndexImpl implements CagraIndex { - private final Dataset dataset; private final CuVSResourcesImpl resources; - private final CagraIndexParams cagraIndexParameters; private final IndexReference cagraIndexReference; private boolean destroyed; @@ -110,11 +108,13 @@ public class CagraIndexImpl implements CagraIndex { * @param dataset the dataset for indexing * @param resources an instance of {@link CuVSResources} */ - private CagraIndexImpl(CagraIndexParams indexParameters, Dataset dataset, CuVSResourcesImpl resources) { - this.cagraIndexParameters = indexParameters; - this.dataset = Objects.requireNonNull(dataset); - this.resources = resources; - this.cagraIndexReference = build(); + private CagraIndexImpl(CagraIndexParams indexParameters, Dataset dataset, CuVSResourcesImpl resources) throws Exception { + Objects.requireNonNull(dataset); + try (dataset) { + this.resources = resources; + assert dataset instanceof DatasetImpl; + this.cagraIndexReference = build(indexParameters, (DatasetImpl) dataset); + } } /** @@ -124,8 +124,6 @@ private CagraIndexImpl(CagraIndexParams indexParameters, Dataset dataset, CuVSRe * @param resources an instance of {@link CuVSResources} */ private CagraIndexImpl(InputStream inputStream, CuVSResourcesImpl resources) throws Throwable { - this.cagraIndexParameters = null; - this.dataset = null; this.resources = resources; this.cagraIndexReference = deserialize(inputStream); } @@ -138,8 +136,6 @@ private CagraIndexImpl(InputStream inputStream, CuVSResourcesImpl resources) thr * @param resources The resources instance */ private CagraIndexImpl(IndexReference indexReference, CuVSResourcesImpl resources) { - this.cagraIndexParameters = null; - this.dataset = null; this.resources = resources; this.cagraIndexReference = indexReference; this.destroyed = false; @@ -155,7 +151,7 @@ private void checkNotDestroyed() { * Invokes the native destroy_cagra_index to de-allocate the CAGRA index */ @Override - public void destroyIndex() throws Throwable { + public void destroyIndex() { checkNotDestroyed(); try (var arena = Arena.ofConfined()) { int returnValue = cuvsCagraIndexDestroy(cagraIndexReference.getMemorySegment()); @@ -163,7 +159,6 @@ public void destroyIndex() throws Throwable { } finally { destroyed = true; } - if (dataset != null) dataset.close(); } /** @@ -173,21 +168,20 @@ public void destroyIndex() throws Throwable { * @return an instance of {@link IndexReference} that holds the pointer to the * index */ - private IndexReference build() { + private IndexReference build(CagraIndexParams indexParameters, DatasetImpl dataset) { try (var localArena = Arena.ofConfined()) { long rows = dataset.size(); long cols = dataset.dimensions(); - MemorySegment indexParamsMemorySegment = cagraIndexParameters != null - ? segmentFromIndexParams(resources, cagraIndexParameters) + MemorySegment indexParamsMemorySegment = indexParameters != null + ? segmentFromIndexParams(resources, indexParameters) : MemorySegment.NULL; - int numWriterThreads = cagraIndexParameters != null ? cagraIndexParameters.getNumWriterThreads() : 1; + int numWriterThreads = indexParameters != null ? indexParameters.getNumWriterThreads() : 1; omp_set_num_threads(numWriterThreads); Arena arena = resources.getArena(); - assert dataset instanceof MemorySegmentProvider; - MemorySegment dataSeg = ((MemorySegmentProvider) dataset).asMemorySegment(); + MemorySegment dataSeg = dataset.asMemorySegment(); long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); MemorySegment stream = arena.allocate(cudaStream_t); @@ -462,16 +456,6 @@ private IndexReference deserialize(InputStream inputStream, int bufferLength) th return indexReference; } - /** - * Gets an instance of {@link CagraIndexParams} - * - * @return an instance of {@link CagraIndexParams} - */ - @Override - public CagraIndexParams getCagraIndexParameters() { - return cagraIndexParameters; - } - /** * Gets an instance of {@link CuVSResources} * diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java index b39987c28b..a6b6798f02 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java @@ -24,7 +24,7 @@ import com.nvidia.cuvs.Dataset; -public class DatasetImpl implements Dataset, MemorySegmentProvider { +public class DatasetImpl implements Dataset { private final Arena arena; protected final MemorySegment seg; private final int size; @@ -38,7 +38,7 @@ public DatasetImpl(int size, int dimensions) { MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout((long)size * dimensions, C_FLOAT); this.arena = Arena.ofShared(); - seg = arena.allocate(dataMemoryLayout); + this.seg = arena.allocate(dataMemoryLayout); } public DatasetImpl(Arena arena, MemorySegment memorySegment, int size, int dimensions) { @@ -57,7 +57,7 @@ public void addVector(float[] vector) { @Override public void close() { - if (arena != null && !arena.scope().isAlive()) { + if (arena != null) { arena.close(); } } @@ -72,7 +72,6 @@ public int dimensions() { return dimensions; } - @Override public MemorySegment asMemorySegment() { return seg; } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/MemorySegmentProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/MemorySegmentProvider.java deleted file mode 100644 index 7ea460a1a8..0000000000 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/MemorySegmentProvider.java +++ /dev/null @@ -1,7 +0,0 @@ -package com.nvidia.cuvs.internal; - -import java.lang.foreign.MemorySegment; - -interface MemorySegmentProvider { - MemorySegment asMemorySegment(); -} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 67da34936e..89339fe41b 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -100,7 +100,6 @@ public Dataset newArrayDataset(float[][] vectors) { int size = vectors.length; int dimensions = vectors[0].length; - // TODO: should we use the Arena from CuVsResources here instead? Arena arena = Arena.ofShared(); var memorySegment = Util.buildMemorySegment(arena, vectors); return new DatasetImpl(arena, memorySegment, size, dimensions); diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java index 2690ab4302..c423da0880 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java @@ -72,6 +72,7 @@ private static void runConcurrently(int nThreads, Supplier runnableSup CompletableFuture.allOf(futures) .exceptionally(t -> { + log.error("Exception while executing runnable", t); fail("Exception while executing runnable: " + t); return null; }) From 840aac512d2575c09ff0f7f52501121465bc3bb1 Mon Sep 17 00:00:00 2001 From: ldematte Date: Fri, 20 Jun 2025 13:40:11 +0200 Subject: [PATCH 08/11] C/CPP changes: add cuvsCagraSerializeToMemory --- cpp/include/cuvs/neighbors/cagra.h | 37 ++++ cpp/src/neighbors/cagra_c.cpp | 45 +++++ .../com/nvidia/cuvs/CagraIndexBenchmarks.java | 183 ------------------ 3 files changed, 82 insertions(+), 183 deletions(-) delete mode 100644 java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index 5959124870..19834d9a29 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -574,6 +574,43 @@ cuvsError_t cuvsCagraSerialize(cuvsResources_t res, cuvsCagraIndex_t index, bool include_dataset); +/** + * @defgroup cagra_c_serialize CAGRA C-API serialize functions + * @{ + */ +/** + * Save the index to file. + * + * Experimental, both the API and the serialization format are subject to change. + * + * @code{.c} + * #include + * + * // Create cuvsResources_t + * cuvsResources_t res; + * cuvsError_t res_create_status = cuvsResourcesCreate(&res); + * + * // create an index with `cuvsCagraBuild` + * size_t length = ...; + * void* buffer = malloc(sizeof(char) * length); + * cuvsCagraSerializeToMemory(res, buffer, length, index, true); + * + * // read buffer + * @endcode + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] buffer pointer to a buffer + * @param[in] size_t buffer_size the size of the buffer + * @param[in] index CAGRA index + * @param[in] include_dataset Whether or not to write out the dataset to the file. + * + */ +cuvsError_t cuvsCagraSerializeToMemory(cuvsResources_t res, + void* buffer, + size_t buffer_size, + cuvsCagraIndex_t index, + bool include_dataset); + /** * Save the CAGRA index to file in hnswlib format. * NOTE: The saved index can only be read by the hnswlib wrapper in cuVS, diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index 656724826e..bcae19e162 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -233,6 +233,30 @@ void _serialize(cuvsResources_t res, cuvs::neighbors::cagra::serialize(*res_ptr, std::string(filename), *index_ptr, include_dataset); } +struct _membuf: std::streambuf +{ + _membuf(char* p, size_t size) + { + setp( p, p + size); + } +}; + +template +void _serialize(cuvsResources_t res, + void* buffer, + size_t buffer_size, + cuvsCagraIndex_t index, + bool include_dataset) +{ + auto res_ptr = reinterpret_cast(res); + auto index_ptr = reinterpret_cast*>(index->addr); + + _membuf sbuf((char*)buffer, buffer_size); + std::ostream out(&sbuf); + + cuvs::neighbors::cagra::serialize(*res_ptr, out, *index_ptr, include_dataset); +} + template void _serialize_to_hnswlib(cuvsResources_t res, const char* filename, cuvsCagraIndex_t index) { @@ -652,6 +676,27 @@ extern "C" cuvsError_t cuvsCagraSerialize(cuvsResources_t res, }); } +extern "C" cuvsError_t cuvsCagraSerializeToMemory(cuvsResources_t res, + void* buffer, + size_t buffer_size, + cuvsCagraIndex_t index, + bool include_dataset) +{ + return cuvs::core::translate_exceptions([=] { + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + _serialize(res, buffer, buffer_size, index, include_dataset); + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + _serialize(res, buffer, buffer_size, index, include_dataset); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + _serialize(res, buffer, buffer_size, index, include_dataset); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + _serialize(res, buffer, buffer_size, index, include_dataset); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); + } + }); +} + extern "C" cuvsError_t cuvsCagraSerializeToHnswlib(cuvsResources_t res, const char* filename, cuvsCagraIndex_t index) diff --git a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java deleted file mode 100644 index 8a0a97936f..0000000000 --- a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java +++ /dev/null @@ -1,183 +0,0 @@ -/* - * Copyright (c) 2025, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.nvidia.cuvs; - -import org.openjdk.jmh.annotations.Benchmark; -import org.openjdk.jmh.annotations.BenchmarkMode; -import org.openjdk.jmh.annotations.Fork; -import org.openjdk.jmh.annotations.Mode; -import org.openjdk.jmh.annotations.OutputTimeUnit; -import org.openjdk.jmh.annotations.Param; -import org.openjdk.jmh.annotations.Scope; -import org.openjdk.jmh.annotations.Setup; -import org.openjdk.jmh.annotations.State; -import org.openjdk.jmh.annotations.TearDown; -import org.openjdk.jmh.infra.Blackhole; - -import java.lang.foreign.*; -import java.nio.file.Files; -import java.nio.file.Path; -import java.util.Random; -import java.util.UUID; -import java.util.concurrent.TimeUnit; - -@Fork(value = 1, warmups = 0) -@State(Scope.Benchmark) -public class CagraIndexBenchmarks { - - @Param({ "1024" }) - private int dims; - - @Param({ "100" }) - private int size; - - private float[][] arrayDataset; - - private Arena arena; - - private MemorySegment memorySegmentDataset; - - private static final Random random = new Random(); - - private static float[][] createSampleData(int size, int dimensions) { - var array = new float[size][dimensions]; - for (int i = 0; i < size; ++i) { - for (int j = 0; j < dimensions; ++j) { - array[i][j] = random.nextFloat(); - } - } - return array; - } - - private static MemorySegment createSampleDataSegment(Arena arena, float[][] array, int size, int dimensions) { - final ValueLayout.OfFloat C_FLOAT = (ValueLayout.OfFloat) Linker.nativeLinker().canonicalLayouts().get("float"); - - MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout((long)size * dimensions, C_FLOAT); - - var segment = arena.allocate(dataMemoryLayout); - for (int i = 0; i < size; ++i) { - var vector = array[i]; - MemorySegment.copy(vector, 0, segment, C_FLOAT, (i * dimensions * C_FLOAT.byteSize()), dimensions); - } - return segment; - } - - @Setup - public void initialize() { - arena = Arena.ofShared(); - arrayDataset = createSampleData(size, dims); - memorySegmentDataset = createSampleDataSegment(arena, arrayDataset, size, dims); - } - - @TearDown - public void cleanUp() { - if (arena != null) { - arena.close(); - } - } - - @Benchmark - public void testIndexingAndSerializeToFile() throws Throwable { - try (CuVSResources resources = CuVSResources.create()) { - // Configure index parameters - CagraIndexParams indexParams = new CagraIndexParams.Builder() - .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT) - .withGraphDegree(1) - .withIntermediateGraphDegree(2) - .withNumWriterThreads(32) - .withMetric(CagraIndexParams.CuvsDistanceType.L2Expanded) - .build(); - - // Create the index with the dataset - CagraIndex index = CagraIndex.newBuilder(resources) - .withDataset(arrayDataset) - .withIndexParams(indexParams) - .build(); - - // Saving the index on to the disk. - var indexFilePath = Path.of(UUID.randomUUID() + ".cag"); - try (var outputStream = Files.newOutputStream(indexFilePath)) { - index.serialize(outputStream); - } - - // Cleanup - Files.deleteIfExists(indexFilePath); - index.destroyIndex(); - } - } - - @Benchmark - public void testIndexingFromHeap(Blackhole blackhole) throws Throwable { - try (CuVSResources resources = CuVSResources.create()) { - - // Configure index parameters - CagraIndexParams indexParams = new CagraIndexParams.Builder() - .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT) - .withGraphDegree(1) - .withIntermediateGraphDegree(2) - .withNumWriterThreads(32) - .withMetric(CagraIndexParams.CuvsDistanceType.L2Expanded) - .build(); - - // Create the index with the dataset - CagraIndex index = CagraIndex.newBuilder(resources) - .withDataset(arrayDataset) - .withIndexParams(indexParams) - .build(); - blackhole.consume(index); - } - } - - @Benchmark - public void testIndexingFromMemorySegment(Blackhole blackhole) throws Throwable { - try (CuVSResources resources = CuVSResources.create()) { - // Configure index parameters - CagraIndexParams indexParams = new CagraIndexParams.Builder() - .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT) - .withGraphDegree(1) - .withIntermediateGraphDegree(2) - .withNumWriterThreads(32) - .withMetric(CagraIndexParams.CuvsDistanceType.L2Expanded) - .build(); - - // Create the index with the dataset - CagraIndex index = CagraIndex.newBuilder(resources) - .withDataset(Dataset.ofMemorySegment(memorySegmentDataset, size, dims)) - .withIndexParams(indexParams) - .build(); - blackhole.consume(index); - } - } - - @Benchmark - @OutputTimeUnit(TimeUnit.NANOSECONDS) - @BenchmarkMode(Mode.AverageTime) - public void testDatasetFromHeap(Blackhole blackhole) throws Throwable { - try (var dataset = Dataset.ofArray(arrayDataset)) { - blackhole.consume(dataset); - } - } - - @Benchmark - @OutputTimeUnit(TimeUnit.NANOSECONDS) - @BenchmarkMode(Mode.AverageTime) - public void testDatasetFromMemorySegment(Blackhole blackhole) throws Throwable { - try (var dataset = Dataset.ofMemorySegment(memorySegmentDataset, size, dims)) { - blackhole.consume(dataset); - } - } -} From 381fb6927989557a2a4d36a473cee8a2bd7f5e37 Mon Sep 17 00:00:00 2001 From: ldematte Date: Fri, 20 Jun 2025 13:40:38 +0200 Subject: [PATCH 09/11] Add serialize(MemorySegment) + test + benchmark --- .../nvidia/cuvs/CagraIndexingBenchmarks.java | 127 ++++++++++++++++++ .../cuvs/CagraSerializationBenchmarks.java | 81 +++++++++++ .../src/main/java/com/nvidia/cuvs/Utils.java | 31 +++++ .../main/java/com/nvidia/cuvs/CagraIndex.java | 9 ++ .../nvidia/cuvs/internal/CagraIndexImpl.java | 15 ++- .../nvidia/cuvs/CagraBuildAndSearchIT.java | 53 ++++++++ 6 files changed, 314 insertions(+), 2 deletions(-) create mode 100644 java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexingBenchmarks.java create mode 100644 java/benchmarks/src/main/java/com/nvidia/cuvs/CagraSerializationBenchmarks.java create mode 100644 java/benchmarks/src/main/java/com/nvidia/cuvs/Utils.java diff --git a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexingBenchmarks.java b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexingBenchmarks.java new file mode 100644 index 0000000000..f1fab12e63 --- /dev/null +++ b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexingBenchmarks.java @@ -0,0 +1,127 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.nvidia.cuvs; + +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.infra.Blackhole; + +import java.lang.foreign.*; +import java.util.concurrent.TimeUnit; + +import static com.nvidia.cuvs.Utils.createSampleData; +import static com.nvidia.cuvs.Utils.createSampleDataSegment; + +@Fork(value = 1, warmups = 0) +@State(Scope.Benchmark) +public class CagraIndexingBenchmarks { + + @Param({ "1024" }) + private int dims; + + @Param({ "100" }) + private int size; + + private float[][] arrayDataset; + + private Arena arena; + + private MemorySegment memorySegmentDataset; + + @Setup + public void initialize() { + arena = Arena.ofShared(); + arrayDataset = createSampleData(size, dims); + memorySegmentDataset = createSampleDataSegment(arena, arrayDataset, size, dims); + } + + @TearDown + public void cleanUp() { + if (arena != null) { + arena.close(); + } + } + + @Benchmark + public void testIndexingFromHeap(Blackhole blackhole) throws Throwable { + try (CuVSResources resources = CuVSResources.create()) { + + // Configure index parameters + CagraIndexParams indexParams = new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(1) + .withIntermediateGraphDegree(2) + .withNumWriterThreads(32) + .withMetric(CagraIndexParams.CuvsDistanceType.L2Expanded) + .build(); + + // Create the index with the dataset + CagraIndex index = CagraIndex.newBuilder(resources) + .withDataset(arrayDataset) + .withIndexParams(indexParams) + .build(); + blackhole.consume(index); + } + } + + @Benchmark + public void testIndexingFromMemorySegment(Blackhole blackhole) throws Throwable { + try (CuVSResources resources = CuVSResources.create()) { + // Configure index parameters + CagraIndexParams indexParams = new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(1) + .withIntermediateGraphDegree(2) + .withNumWriterThreads(32) + .withMetric(CagraIndexParams.CuvsDistanceType.L2Expanded) + .build(); + + // Create the index with the dataset + CagraIndex index = CagraIndex.newBuilder(resources) + .withDataset(Dataset.ofMemorySegment(memorySegmentDataset, size, dims)) + .withIndexParams(indexParams) + .build(); + blackhole.consume(index); + } + } + + @Benchmark + @OutputTimeUnit(TimeUnit.NANOSECONDS) + @BenchmarkMode(Mode.AverageTime) + public void testDatasetFromHeap(Blackhole blackhole) throws Throwable { + try (var dataset = Dataset.ofArray(arrayDataset)) { + blackhole.consume(dataset); + } + } + + @Benchmark + @OutputTimeUnit(TimeUnit.NANOSECONDS) + @BenchmarkMode(Mode.AverageTime) + public void testDatasetFromMemorySegment(Blackhole blackhole) throws Throwable { + try (var dataset = Dataset.ofMemorySegment(memorySegmentDataset, size, dims)) { + blackhole.consume(dataset); + } + } +} diff --git a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraSerializationBenchmarks.java b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraSerializationBenchmarks.java new file mode 100644 index 0000000000..d238a9d3e2 --- /dev/null +++ b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraSerializationBenchmarks.java @@ -0,0 +1,81 @@ +package com.nvidia.cuvs; + +import org.openjdk.jmh.annotations.*; + +import java.lang.foreign.Arena; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.UUID; + +import static com.nvidia.cuvs.Utils.createSampleData; + +@Fork(value = 1, warmups = 0) +@State(Scope.Benchmark) +public class CagraSerializationBenchmarks { + + @Param({ "1024" }) + private int dims; + + @Param({ "100" }) + private int size; + + private Arena arena; + + private CuVSResources resources; + + private CagraIndex index; + + @Setup + public void initialize() throws Throwable { + resources = CuVSResources.create(); + arena = Arena.ofShared(); + var arrayDataset = createSampleData(size, dims); + + // Configure index parameters + CagraIndexParams indexParams = new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(1) + .withIntermediateGraphDegree(2) + .withNumWriterThreads(32) + .withMetric(CagraIndexParams.CuvsDistanceType.L2Expanded) + .build(); + + // Create the index with the dataset + index = CagraIndex.newBuilder(resources) + .withDataset(arrayDataset) + .withIndexParams(indexParams) + .build(); + } + + @TearDown + public void cleanUp() throws Throwable { + if (resources != null) { + resources.close(); + } + if (arena != null) { + arena.close(); + } + if (index != null) { + index.destroyIndex(); + } + } + + @Benchmark + public void testSerializeToFile() throws Throwable { + var indexFilePath = Path.of(UUID.randomUUID() + ".cag"); + try (var outputStream = Files.newOutputStream(indexFilePath)) { + index.serialize(outputStream); + } + Files.deleteIfExists(indexFilePath); + } + + @Benchmark + public void testSerializeToMemory() throws Throwable { + try (var arena = Arena.ofConfined()) { + var buffer = arena.allocate(1024 * 1024); + index.serialize((Object) buffer); + // TODO: reinterpret should happen in serialize, with data from the C call + + } + } +} diff --git a/java/benchmarks/src/main/java/com/nvidia/cuvs/Utils.java b/java/benchmarks/src/main/java/com/nvidia/cuvs/Utils.java new file mode 100644 index 0000000000..4a0c1200e6 --- /dev/null +++ b/java/benchmarks/src/main/java/com/nvidia/cuvs/Utils.java @@ -0,0 +1,31 @@ +package com.nvidia.cuvs; + +import java.lang.foreign.*; +import java.util.Random; + +class Utils { + private static final Random random = new Random(); + + static float[][] createSampleData(int size, int dimensions) { + var array = new float[size][dimensions]; + for (int i = 0; i < size; ++i) { + for (int j = 0; j < dimensions; ++j) { + array[i][j] = random.nextFloat(); + } + } + return array; + } + + static MemorySegment createSampleDataSegment(Arena arena, float[][] array, int size, int dimensions) { + final ValueLayout.OfFloat C_FLOAT = (ValueLayout.OfFloat) Linker.nativeLinker().canonicalLayouts().get("float"); + + MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout((long)size * dimensions, C_FLOAT); + + var segment = arena.allocate(dataMemoryLayout); + for (int i = 0; i < size; ++i) { + var vector = array[i]; + MemorySegment.copy(vector, 0, segment, C_FLOAT, (i * dimensions * C_FLOAT.byteSize()), dimensions); + } + return segment; + } +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index aae46dcb7a..b786515d2c 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java @@ -99,6 +99,15 @@ default void serialize(OutputStream outputStream, Path tempFile) throws Throwabl */ void serialize(OutputStream outputStream, Path tempFile, int bufferLength) throws Throwable; + /** + * A method to persist a CAGRA index using a {@link java.lang.foreign.MemorySegment} + * for writing index bytes. + * + * @param memoryStream an instance of {@link java.lang.foreign.MemorySegment} to write the index + * bytes into + */ + void serialize(Object memoryStream); + /** * A method to create and persist HNSW index from CAGRA index using an instance * of {@link OutputStream} and path to the intermediate temporary file. diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 8563a998e4..e61454563f 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -27,24 +27,25 @@ import static com.nvidia.cuvs.internal.common.Util.cudaMemcpy; import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.*; import static com.nvidia.cuvs.internal.common.Util.prepareTensor; +import static com.nvidia.cuvs.internal.panama.headers_h.cudaStream_t; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraBuild; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraDeserialize; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraIndexCreate; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraIndexDestroy; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraIndexGetDims; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraIndex_t; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraSearch; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraMerge; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraMergeParams_t; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraSearch; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraSerialize; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraSerializeToHnswlib; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsCagraSerializeToMemory; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsRMMAlloc; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsRMMFree; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsResources_t; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamGet; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamSync; import static com.nvidia.cuvs.internal.panama.headers_h.omp_set_num_threads; -import static com.nvidia.cuvs.internal.panama.headers_h.cudaStream_t; import java.io.FileInputStream; import java.io.FileOutputStream; @@ -349,6 +350,16 @@ public SearchResults search(CagraQuery query) throws Throwable { } } + public void serialize(Object memorySegment) { + assert memorySegment instanceof MemorySegment; + var buffer = (MemorySegment) memorySegment; + checkNotDestroyed(); + long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); + var returnValue = cuvsCagraSerializeToMemory(cuvsRes, buffer, buffer.byteSize(), + cagraIndexReference.getMemorySegment(), true); + checkCuVSError(returnValue, "cuvsCagraSerializeToMemory"); + } + @Override public void serialize(OutputStream outputStream) throws Throwable { Path path = Files.createTempFile(resources.tempDirectory(), UUID.randomUUID().toString(), ".cag"); diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java index c423da0880..97b985e2ee 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java @@ -17,6 +17,7 @@ package com.nvidia.cuvs; import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.fail; @@ -25,7 +26,11 @@ import java.io.FileInputStream; import java.io.FileOutputStream; import java.io.InputStream; +import java.lang.foreign.Arena; +import java.lang.foreign.ValueLayout; import java.lang.invoke.MethodHandles; +import java.nio.file.Files; +import java.nio.file.Path; import java.util.Arrays; import java.util.BitSet; import java.util.List; @@ -237,6 +242,54 @@ public void testPrefilteringReducesResults() throws Throwable { } } + @Test + public void testIndexSerialization() throws Throwable { + float[][] dataset = createSampleData(); + + try (CuVSResources resources = CuVSResources.create()) { + + // Configure index parameters + CagraIndexParams indexParams = new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(1) + .withIntermediateGraphDegree(2) + .withNumWriterThreads(32) + .withMetric(CuvsDistanceType.L2Expanded) + .build(); + + // Create the index with the dataset + CagraIndex index = CagraIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(indexParams) + .build(); + + var indexFile = Path.of(UUID.randomUUID() + ".cag"); + try { + // Saving the index on to the disk. + try (var output = Files.newOutputStream(indexFile)) { + index.serialize(output); + } + + var indexFromFile = Files.readAllBytes(indexFile); + byte[] indexFromMemory; + try (var arena = Arena.ofConfined()) { + var buffer = arena.allocate(1024 * 1024); + index.serialize(buffer); + // TODO: reinterpret should happen in serialize, with data from the C call + indexFromMemory = buffer.reinterpret(indexFromFile.length).toArray(ValueLayout.JAVA_BYTE); + } + + assertNotNull(indexFromFile); + assertNotNull(indexFromMemory); + assertArrayEquals(indexFromFile, indexFromMemory); + + } finally { + Files.deleteIfExists(indexFile); + index.destroyIndex(); + } + } + } + private void indexAndQueryOnce(float[][] dataset, List map, float[][] queries, List> expectedResults, CuVSResources resources) { From dfce2c2ea17e53387140ce3cfad499bfba8e7a85 Mon Sep 17 00:00:00 2001 From: ldematte Date: Fri, 20 Jun 2025 18:26:29 +0200 Subject: [PATCH 10/11] Update API to report and use the number of bytes written --- cpp/include/cuvs/neighbors/cagra.h | 4 ++-- cpp/src/neighbors/cagra_c.cpp | 21 ++++++++++++------- .../cuvs/CagraSerializationBenchmarks.java | 2 -- .../main/java/com/nvidia/cuvs/CagraIndex.java | 3 ++- .../nvidia/cuvs/internal/CagraIndexImpl.java | 13 ++++++++---- .../nvidia/cuvs/CagraBuildAndSearchIT.java | 6 +++--- 6 files changed, 29 insertions(+), 20 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index 19834d9a29..561f4a833f 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -600,14 +600,14 @@ cuvsError_t cuvsCagraSerialize(cuvsResources_t res, * * @param[in] res cuvsResources_t opaque C handle * @param[in] buffer pointer to a buffer - * @param[in] size_t buffer_size the size of the buffer + * @param[in/out] length the size of the buffer; the function will update it with the number of bytes written * @param[in] index CAGRA index * @param[in] include_dataset Whether or not to write out the dataset to the file. * */ cuvsError_t cuvsCagraSerializeToMemory(cuvsResources_t res, void* buffer, - size_t buffer_size, + size_t* length, cuvsCagraIndex_t index, bool include_dataset); diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index bcae19e162..9e76addb53 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -239,22 +239,27 @@ struct _membuf: std::streambuf { setp( p, p + size); } + size_t written() + { + return pptr()-pbase(); + } }; template void _serialize(cuvsResources_t res, void* buffer, - size_t buffer_size, + size_t* length, cuvsCagraIndex_t index, bool include_dataset) { auto res_ptr = reinterpret_cast(res); auto index_ptr = reinterpret_cast*>(index->addr); - _membuf sbuf((char*)buffer, buffer_size); - std::ostream out(&sbuf); + _membuf stream_buffer((char*)buffer, *length); + std::ostream out(&stream_buffer); cuvs::neighbors::cagra::serialize(*res_ptr, out, *index_ptr, include_dataset); + *length = stream_buffer.written(); } template @@ -678,19 +683,19 @@ extern "C" cuvsError_t cuvsCagraSerialize(cuvsResources_t res, extern "C" cuvsError_t cuvsCagraSerializeToMemory(cuvsResources_t res, void* buffer, - size_t buffer_size, + size_t* length, cuvsCagraIndex_t index, bool include_dataset) { return cuvs::core::translate_exceptions([=] { if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { - _serialize(res, buffer, buffer_size, index, include_dataset); + _serialize(res, buffer, length, index, include_dataset); } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { - _serialize(res, buffer, buffer_size, index, include_dataset); + _serialize(res, buffer, length, index, include_dataset); } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { - _serialize(res, buffer, buffer_size, index, include_dataset); + _serialize(res, buffer, length, index, include_dataset); } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { - _serialize(res, buffer, buffer_size, index, include_dataset); + _serialize(res, buffer, length, index, include_dataset); } else { RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); } diff --git a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraSerializationBenchmarks.java b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraSerializationBenchmarks.java index d238a9d3e2..2056f56c14 100644 --- a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraSerializationBenchmarks.java +++ b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraSerializationBenchmarks.java @@ -74,8 +74,6 @@ public void testSerializeToMemory() throws Throwable { try (var arena = Arena.ofConfined()) { var buffer = arena.allocate(1024 * 1024); index.serialize((Object) buffer); - // TODO: reinterpret should happen in serialize, with data from the C call - } } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index b786515d2c..38b7e08f70 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java @@ -105,8 +105,9 @@ default void serialize(OutputStream outputStream, Path tempFile) throws Throwabl * * @param memoryStream an instance of {@link java.lang.foreign.MemorySegment} to write the index * bytes into + * @return a MemorySegment containing the serialized data */ - void serialize(Object memoryStream); + Object serialize(Object memoryStream); /** * A method to create and persist HNSW index from CAGRA index using an instance diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index e61454563f..3b8af37425 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -350,14 +350,19 @@ public SearchResults search(CagraQuery query) throws Throwable { } } - public void serialize(Object memorySegment) { + @Override + public Object serialize(Object memorySegment) { assert memorySegment instanceof MemorySegment; var buffer = (MemorySegment) memorySegment; checkNotDestroyed(); long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); - var returnValue = cuvsCagraSerializeToMemory(cuvsRes, buffer, buffer.byteSize(), - cagraIndexReference.getMemorySegment(), true); - checkCuVSError(returnValue, "cuvsCagraSerializeToMemory"); + try (var arena = Arena.ofConfined()) { + var length = arena.allocateFrom(ValueLayout.JAVA_LONG, buffer.byteSize()); + var returnValue = cuvsCagraSerializeToMemory(cuvsRes, buffer, length, + cagraIndexReference.getMemorySegment(), true); + checkCuVSError(returnValue, "cuvsCagraSerializeToMemory"); + return buffer.reinterpret(length.get(ValueLayout.JAVA_LONG, 0)); + } } @Override diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java index 97b985e2ee..addb051890 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java @@ -27,6 +27,7 @@ import java.io.FileOutputStream; import java.io.InputStream; import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; import java.lang.foreign.ValueLayout; import java.lang.invoke.MethodHandles; import java.nio.file.Files; @@ -274,9 +275,8 @@ public void testIndexSerialization() throws Throwable { byte[] indexFromMemory; try (var arena = Arena.ofConfined()) { var buffer = arena.allocate(1024 * 1024); - index.serialize(buffer); - // TODO: reinterpret should happen in serialize, with data from the C call - indexFromMemory = buffer.reinterpret(indexFromFile.length).toArray(ValueLayout.JAVA_BYTE); + var serialized = (MemorySegment)index.serialize((Object) buffer); + indexFromMemory = serialized.toArray(ValueLayout.JAVA_BYTE); } assertNotNull(indexFromFile); From f8500a42c46f30ca1701ac0256abb5e029aad62d Mon Sep 17 00:00:00 2001 From: ldematte Date: Thu, 3 Jul 2025 14:56:16 +0200 Subject: [PATCH 11/11] Fix formatting --- cpp/include/cuvs/neighbors/cagra.h | 12 +++++------- cpp/src/neighbors/cagra_c.cpp | 27 +++++++-------------------- java/.gitignore | 2 -- java/benchmarks/README.md | 2 +- 4 files changed, 13 insertions(+), 30 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index 561f4a833f..00da741932 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -600,16 +600,14 @@ cuvsError_t cuvsCagraSerialize(cuvsResources_t res, * * @param[in] res cuvsResources_t opaque C handle * @param[in] buffer pointer to a buffer - * @param[in/out] length the size of the buffer; the function will update it with the number of bytes written + * @param[in/out] length the size of the buffer; the function will update it with the number of + * bytes written * @param[in] index CAGRA index * @param[in] include_dataset Whether or not to write out the dataset to the file. * */ -cuvsError_t cuvsCagraSerializeToMemory(cuvsResources_t res, - void* buffer, - size_t* length, - cuvsCagraIndex_t index, - bool include_dataset); +cuvsError_t cuvsCagraSerializeToMemory( + cuvsResources_t res, void* buffer, size_t* length, cuvsCagraIndex_t index, bool include_dataset); /** * Save the CAGRA index to file in hnswlib format. diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index 9e76addb53..8d3deb4885 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -233,24 +233,14 @@ void _serialize(cuvsResources_t res, cuvs::neighbors::cagra::serialize(*res_ptr, std::string(filename), *index_ptr, include_dataset); } -struct _membuf: std::streambuf -{ - _membuf(char* p, size_t size) - { - setp( p, p + size); - } - size_t written() - { - return pptr()-pbase(); - } +struct _membuf : std::streambuf { + _membuf(char* p, size_t size) { setp(p, p + size); } + size_t written() { return pptr() - pbase(); } }; template -void _serialize(cuvsResources_t res, - void* buffer, - size_t* length, - cuvsCagraIndex_t index, - bool include_dataset) +void _serialize( + cuvsResources_t res, void* buffer, size_t* length, cuvsCagraIndex_t index, bool include_dataset) { auto res_ptr = reinterpret_cast(res); auto index_ptr = reinterpret_cast*>(index->addr); @@ -681,11 +671,8 @@ extern "C" cuvsError_t cuvsCagraSerialize(cuvsResources_t res, }); } -extern "C" cuvsError_t cuvsCagraSerializeToMemory(cuvsResources_t res, - void* buffer, - size_t* length, - cuvsCagraIndex_t index, - bool include_dataset) +extern "C" cuvsError_t cuvsCagraSerializeToMemory( + cuvsResources_t res, void* buffer, size_t* length, cuvsCagraIndex_t index, bool include_dataset) { return cuvs::core::translate_exceptions([=] { if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { diff --git a/java/.gitignore b/java/.gitignore index ab06477b43..a240aca207 100644 --- a/java/.gitignore +++ b/java/.gitignore @@ -6,5 +6,3 @@ openjdk-22-jextract* /cuvs-java/bin/ /cuvs-java/src/main/java22/com/nvidia/cuvs/internal/panama/ /cuvs-java/*.cag - - diff --git a/java/benchmarks/README.md b/java/benchmarks/README.md index 447c578448..1740e52811 100644 --- a/java/benchmarks/README.md +++ b/java/benchmarks/README.md @@ -24,4 +24,4 @@ It is possible to change the dataset size and the vectors dimension via 2 parame ```shell java -jar target/benchmarks.jar -p size=4 -p dims=4 ``` -Use `java -jar target/benchmarks.jar -h` for details on the options to fine-tune your benchmark runs. \ No newline at end of file +Use `java -jar target/benchmarks.jar -h` for details on the options to fine-tune your benchmark runs.