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 51403982dd..96e31431f8 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 @@ -336,6 +336,23 @@ static CagraIndex merge(CagraIndex[] indexes, CagraIndexParams mergeParams) thro return CuVSProvider.provider().mergeCagraIndexes(indexes, mergeParams); } + /** + * Reports whether the rows of {@code dataset} already sit at the row stride CAGRA requires, which + * is the row length in bytes rounded up to a 16 byte boundary. + * + *

Use it to pick between the two padded dataset factories: a matrix that is already padded has + * to go through {@link #makePaddedDatasetView(CuVSMatrix)}, because cuVS rejects a request to + * copy it into padded storage it already occupies, and one that is not has to go through + * {@link #makePaddedDataset(CuVSMatrix)}. + * + * @param dataset the matrix to inspect + * @return true when the rows are already padded the way CAGRA requires + */ + static boolean isPaddedDataset(CuVSMatrix dataset) { + Objects.requireNonNull(dataset); + return CuVSProvider.provider().isCagraPaddedDataset(dataset); + } + /** * Builder helps configure and create an instance of {@link CagraIndex}. */ 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 05bbc4ab0f..670f44eda6 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 @@ -188,6 +188,25 @@ default CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraIndexParams merg return mergeCagraIndexes(indexes); } + /** + * Reports whether the rows of {@code dataset} already sit at the row stride CAGRA requires, which + * is the row length in bytes rounded up to a 16 byte boundary. + * + *

This is the question that decides which of the two padded dataset factories a caller has to + * use: {@link CagraIndex#makePaddedDatasetView(CuVSMatrix)} for a device matrix that is already at + * that stride, and {@link CagraIndex#makePaddedDataset(CuVSMatrix)} for one that is not. Asking + * for the wrong one is an error rather than an inefficiency, and the stride of a matrix is not + * visible outside this library, so callers cannot answer it for themselves. + * + * @param dataset the matrix to inspect + * @return true when the rows are already padded the way CAGRA requires + * @throws UnsupportedOperationException if this provider cannot answer + */ + default boolean isCagraPaddedDataset(CuVSMatrix dataset) { + throw new UnsupportedOperationException( + "Padded layout detection is not supported by " + getClass().getName()); + } + /** * Creates a device-backed multi-partition filter handle from the pre-packed combined bitset. * Per-partition bit offsets are recomputed inside cuVS from the index sizes. 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 2a9f9ac738..66b0b8cc47 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 @@ -85,6 +85,11 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) { throw new UnsupportedOperationException(reasons); } + @Override + public boolean isCagraPaddedDataset(CuVSMatrix dataset) { + throw new UnsupportedOperationException(reasons); + } + @Override public CuVSMatrix.Builder newHostMatrixBuilder( long size, long dimensions, CuVSMatrix.DataType dataType) { 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 691c99e93f..9506b130cd 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 @@ -263,6 +263,19 @@ private static int elementSizeBytes(CuVSMatrix.DataType dataType) { }; } + /** + * True when the matrix row width matches CAGRA's required padded width for its logical column + * count and element type. The stride lives on the internal matrix type, so this is the only place + * that can answer the question; {@link com.nvidia.cuvs.CagraIndex#isPaddedDataset} routes here + * through the provider. + */ + public static boolean isPaddedDataset(CuVSMatrix dataset) { + if (!(dataset instanceof CuVSMatrixInternal datasetInternal)) { + throw new IllegalArgumentException("dataset must be a CuVSMatrixInternal matrix"); + } + return isCagraPaddedLayout(datasetInternal); + } + /** * True when the matrix row width matches CAGRA's required padded width for its * logical column count and element type. 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 4d1dbecf0f..9cc4a5499c 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 @@ -317,6 +317,11 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraIndexParams merge return CagraIndexImpl.merge(indexes, mergeParams); } + @Override + public boolean isCagraPaddedDataset(CuVSMatrix dataset) { + return CagraIndexImpl.isPaddedDataset(dataset); + } + @Override public GPUInfoProvider gpuInfoProvider() { return new GPUInfoProviderImpl(); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index 97fedf924f..aead683dc4 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -206,6 +206,14 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro writeCagraIndex(cagraIndexOutputStream, cagraDataset); } catch (Throwable t) { // Fallback to brute force in a few cases, for now. + // Log it to make it more obvious that this is what is happening. + info( + infoStream, + COMPONENT, + "CAGRA build failed for field \"" + + fieldInfo.name + + "\", falling back to a brute force index: " + + t); Utils.handleThrowableWithIgnore(t, t.getMessage()); indexType = IndexType.BRUTE_FORCE; } @@ -248,10 +256,23 @@ private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwab .withDataset(dataset) .withIndexParams(params) .build(); - var deviceVectors = dataset.toDevice(getCuVSResourcesInstance()); - var indexDataset = index.makePaddedDataset(deviceVectors)) { - index.updateDataset(indexDataset); - index.serialize(os); + var deviceVectors = dataset.toDevice(getCuVSResourcesInstance())) { + /* + * cuVS rejects makePaddedDataset for a device matrix whose rows already sit at the required + * stride, and asks for a view over that storage instead. Copying would be pointless there + * anyway, so pick the factory that matches the layout. + */ + if (CagraIndex.isPaddedDataset(deviceVectors)) { + try (var indexDatasetView = index.makePaddedDatasetView(deviceVectors)) { + index.updateDataset(indexDatasetView); + index.serialize(os); + } + } else { + try (var indexDataset = index.makePaddedDataset(deviceVectors)) { + index.updateDataset(indexDataset); + index.serialize(os); + } + } } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java index b24d35ffa7..1c74649843 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java @@ -85,6 +85,11 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] arg0) throws Throwable { return delegate.mergeCagraIndexes(arg0); } + @Override + public boolean isCagraPaddedDataset(CuVSMatrix arg0) { + return delegate.isCagraPaddedDataset(arg0); + } + @Override public GPUInfoProvider gpuInfoProvider() { return delegate.gpuInfoProvider(); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexAtAlignedDimensions.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexAtAlignedDimensions.java new file mode 100644 index 0000000000..89d58aa682 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexAtAlignedDimensions.java @@ -0,0 +1,115 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; + +import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.InfoStream; +import org.junit.Assume; +import org.junit.Test; + +/** + * A CAGRA row is padded to a 16 byte boundary, so a device matrix whose dimension is already a + * multiple of four sits at the required stride. cuVS refuses to build an owning padded copy of such + * a matrix and asks for a view instead, and the writer swallows a failed CAGRA build by falling + * back to a brute force index. The two together are silent: search keeps returning correct results + * while nothing on the GPU is a CAGRA index any more. + * + *

These tests pin the dimensions on both sides of that boundary. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestCagraIndexAtAlignedDimensions extends LuceneTestCase { + + @Test + public void testCagraIsBuiltAtAnAlignedDimension() throws IOException { + // 128 floats is 512 bytes, an exact multiple of the 16 byte CAGRA row alignment. + assertCagraIsBuilt(128); + } + + @Test + public void testCagraIsBuiltAtAnUnalignedDimension() throws IOException { + // 127 floats is not, so the writer has to fall back to an owning padded copy. + assertCagraIsBuilt(127); + } + + /** Indexes a segment of the given dimension and fails if the CAGRA build did not survive it. */ + private void assertCagraIsBuilt(int dimension) throws IOException { + Assume.assumeTrue("Requires a GPU", isSupported()); + + RecordingInfoStream infoStream = new RecordingInfoStream(); + try (Directory directory = newDirectory()) { + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec( + TestUtil.alwaysKnnVectorsFormat( + new CuVS2510GPUVectorsFormat( + new GPUSearchParams.Builder().withIndexType(IndexType.CAGRA).build()))) + .setInfoStream(infoStream); + + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (int i = 0; i < 64; i++) { + float[] vector = new float[dimension]; + for (int d = 0; d < dimension; d++) { + vector[d] = random().nextFloat(); + } + Document doc = new Document(); + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + } + writer.commit(); + } + } + + assertTrue( + "The CAGRA build fell back to brute force at dimension " + + dimension + + ", messages: " + + infoStream.messages(), + infoStream.cagraBuildFailures().isEmpty()); + } + + /** An InfoStream that keeps the messages, so that a test can tell which index type was built. */ + private static class RecordingInfoStream extends InfoStream { + + private final List messages = Collections.synchronizedList(new ArrayList<>()); + + @Override + public void message(String component, String message) { + messages.add(component + ": " + message); + } + + @Override + public boolean isEnabled(String component) { + return true; + } + + @Override + public void close() {} + + List messages() { + synchronized (messages) { + return List.copyOf(messages); + } + } + + List cagraBuildFailures() { + return messages().stream().filter(message -> message.contains("CAGRA build failed")).toList(); + } + } +}