diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindexparams.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindexparams.md index d1e8bccfc3..8a7a98b629 100644 --- a/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindexparams.md +++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindexparams.md @@ -508,17 +508,21 @@ public Builder withNumWriterThreads(int numWriterThreads) Sets the number of writer threads to use for indexing. +This sets the OpenMP thread count for the native build or merge on the calling thread, and +the previous count is restored afterwards. A value of zero or less leaves the OpenMP setting +untouched, so the native code uses the OpenMP default (normally one thread per core). + **Parameters** | Name | Description | | --- | --- | -| `numWriterThreads` | number of writer threads to use | +| `numWriterThreads` | number of writer threads to use, or zero or less for the OpenMP default | **Returns** an instance of Builder -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java:532`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java:536`_ ### withCuVSIvfPqParams @@ -538,7 +542,7 @@ Sets the IVF_PQ index parameters. an instance of Builder -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java:543`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java:547`_ ### withCuVSAceParams @@ -558,7 +562,7 @@ Sets the ACE index parameters. an instance of Builder -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java:554`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java:558`_ ### build @@ -572,6 +576,6 @@ Builds an instance of `CagraIndexParams`. an instance of `CagraIndexParams` -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java:564`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java:568`_ _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java:18`_ diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java index b12fdbc11f..65ba9054fb 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java @@ -527,7 +527,11 @@ public Builder withNNDescentNumIterations(long nnDescentNiter) { /** * Sets the number of writer threads to use for indexing. * - * @param numWriterThreads number of writer threads to use + *

This sets the OpenMP thread count for the native build or merge on the calling thread, and + * the previous count is restored afterwards. A value of zero or less leaves the OpenMP setting + * untouched, so the native code uses the OpenMP default (normally one thread per core). + * + * @param numWriterThreads number of writer threads to use, or zero or less for the OpenMP default * @return an instance of Builder */ public Builder withNumWriterThreads(int numWriterThreads) { 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 a494b72e36..18ecb26309 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 @@ -119,8 +119,6 @@ private IndexReference build( MemorySegment datasetMemSegment = dataset.memorySegment(); - omp_set_num_threads(bruteForceIndexParams.getNumWriterThreads()); - long datasetBytes = C_FLOAT_BYTE_SIZE * rows * cols; var index = createBruteForceIndex(); @@ -152,8 +150,6 @@ private IndexReference build( tensorDataArena, index); } - } finally { - omp_set_num_threads(1); } } 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 bac4a7f2c3..9a31038c79 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 @@ -23,6 +23,7 @@ import com.nvidia.cuvs.internal.common.CloseableHandle; import com.nvidia.cuvs.internal.common.CloseableRMMAllocation; import com.nvidia.cuvs.internal.common.CompositeCloseableHandle; +import com.nvidia.cuvs.internal.common.OmpThreadScope; import com.nvidia.cuvs.internal.common.Util; import com.nvidia.cuvs.internal.panama.*; import java.io.FileInputStream; @@ -179,19 +180,24 @@ private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixInterna var resourcesAccessor = resources.access()) { MemorySegment indexParamsMemorySegment = indexParams.handle(); - int numWriterThreads = indexParameters != null ? indexParameters.getNumWriterThreads() : 1; + // Without parameters, leave the OpenMP setting alone so native code uses its default. + int numWriterThreads = indexParameters != null ? indexParameters.getNumWriterThreads() : 0; MemorySegment datasetView = MemorySegment.NULL; boolean success = false; var index = createCagraIndex(); - try { - omp_set_num_threads(numWriterThreads); - - if (cuvsCagraIndexParams.build_algo(indexParamsMemorySegment) - == 1) { // when build algo is IVF_PQ - MemorySegment cuvsIvfPqIndexParamsMS = - cuvsIvfPqParams.ivf_pq_build_params( - cuvsCagraIndexParams.graph_build_params(indexParamsMemorySegment)); + try (var ompThreads = OmpThreadScope.of(numWriterThreads)) { + // With default parameters (no params given) the IVF-PQ build params are not allocated + // (NULL); the C layer then derives them from the dataset, so there is nothing to trim. + MemorySegment cuvsIvfPqIndexParamsMS = MemorySegment.NULL; + if (cuvsCagraIndexParams.build_algo(indexParamsMemorySegment) == 1) { // IVF_PQ + MemorySegment graphBuildParams = + cuvsCagraIndexParams.graph_build_params(indexParamsMemorySegment); + if (graphBuildParams.address() != 0) { + cuvsIvfPqIndexParamsMS = cuvsIvfPqParams.ivf_pq_build_params(graphBuildParams); + } + } + if (cuvsIvfPqIndexParamsMS.address() != 0) { int n_lists = cuvsIvfPqIndexParams.n_lists(cuvsIvfPqIndexParamsMS); // As rows cannot be less than n_lists value so trim down. cuvsIvfPqIndexParams.n_lists( @@ -225,7 +231,6 @@ private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixInterna if (!success) { quietly(() -> checkCuVSError(cuvsCagraIndexDestroy(index), "cuvsCagraIndexDestroy")); } - quietly(() -> omp_set_num_threads(1)); } } } @@ -246,14 +251,13 @@ private IndexReference buildBbq( var quantizerHandles = new ArrayList(bbqQuantizers.length); MemorySegment bbqDataset = MemorySegment.NULL; MemorySegment paddedDataset = MemorySegment.NULL; - int numWriterThreads = indexParameters != null ? indexParameters.getNumWriterThreads() : 1; + int numWriterThreads = indexParameters != null ? indexParameters.getNumWriterThreads() : 0; boolean success = false; var index = createCagraIndex(); - try { - // nn-descent, which a BBQ build always uses, does a substantial amount of its work on - // the host under OpenMP. - omp_set_num_threads(numWriterThreads); + // nn-descent, which a BBQ build always uses, does a substantial amount of its work on + // the host under OpenMP. + try (var ompThreads = OmpThreadScope.of(numWriterThreads)) { for (BbqQuantizer quantizer : bbqQuantizers) { var quantizerPtr = localArena.allocate(cuvsBbqQuantizer_t); checkCuVSError( @@ -328,7 +332,6 @@ private IndexReference buildBbq( quietly( () -> checkCuVSError(cuvsBbqQuantizerDestroy(quantizer), "cuvsBbqQuantizerDestroy")); } - quietly(() -> omp_set_num_threads(1)); } } } @@ -1241,8 +1244,10 @@ public static CagraIndex merge( var mergedIndex = createCagraIndex(); CagraIndexImpl merged = null; + int numWriterThreads = mergeParams != null ? mergeParams.getNumWriterThreads() : 0; try (var nativeMergeParams = segmentFromIndexParams(mergeParams); - var resourcesAccessor = resources.access()) { + var resourcesAccessor = resources.access(); + var ompThreads = OmpThreadScope.of(numWriterThreads)) { var cuvsRes = resourcesAccessor.handle(); // The words the merge filter points at have to outlive the merge call, so the diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/OmpThreadScope.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/OmpThreadScope.java new file mode 100644 index 0000000000..7b0422d5ef --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/OmpThreadScope.java @@ -0,0 +1,45 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.internal.common; + +import static com.nvidia.cuvs.internal.panama.headers_h.omp_get_max_threads; +import static com.nvidia.cuvs.internal.panama.headers_h.omp_set_num_threads; + +/** + * Sets the OpenMP thread count for the native work done by the current thread, and restores the + * previous value on {@link #close()}. + * + *

OpenMP's {@code nthreads-var} is per thread: {@code omp_set_num_threads} only affects parallel + * regions started by the calling thread, so concurrent builds on other threads keep their own + * values. What must not happen is leaving this thread's value changed after the call returns, which + * is why the previous value is restored rather than reset to a fixed number. + */ +public final class OmpThreadScope implements AutoCloseable { + private final int previous; + + private OmpThreadScope(int previous) { + this.previous = previous; + } + + /** + * Applies {@code numThreads} for the current thread until the returned scope is closed. A value + * of zero or less leaves the current setting untouched, so native code uses the OpenMP default. + */ + public static OmpThreadScope of(int numThreads) { + if (numThreads <= 0) { + return new OmpThreadScope(0); + } + int previous = omp_get_max_threads(); + omp_set_num_threads(numThreads); + return new OmpThreadScope(previous); + } + + @Override + public void close() { + if (previous > 0) { + omp_set_num_threads(previous); + } + } +} diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/OmpThreadsIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/OmpThreadsIT.java new file mode 100644 index 0000000000..eb855a27b5 --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/OmpThreadsIT.java @@ -0,0 +1,125 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs; + +import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static org.junit.Assert.assertEquals; + +import com.carrotsearch.randomizedtesting.RandomizedRunner; +import java.lang.foreign.Arena; +import java.lang.foreign.FunctionDescriptor; +import java.lang.foreign.Linker; +import java.lang.foreign.SymbolLookup; +import java.lang.foreign.ValueLayout; +import java.lang.invoke.MethodHandle; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; + +/** A build must not leave the calling thread's OpenMP thread count changed. */ +@RunWith(RandomizedRunner.class) +public class OmpThreadsIT extends CuVSTestCase { + + // Differs from both the old hard-coded reset (1) and the numWriterThreads the tests request (3). + private static final int BASELINE = 5; + + private MethodHandle ompGetMaxThreads; + private MethodHandle ompSetNumThreads; + private int original; + + @Before + public void setup() throws Throwable { + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); + initializeRandom(); + var libgomp = SymbolLookup.libraryLookup("libgomp.so.1", Arena.global()); + var linker = Linker.nativeLinker(); + ompGetMaxThreads = + linker.downcallHandle( + libgomp.find("omp_get_max_threads").orElseThrow(), + FunctionDescriptor.of(ValueLayout.JAVA_INT)); + ompSetNumThreads = + linker.downcallHandle( + libgomp.find("omp_set_num_threads").orElseThrow(), + FunctionDescriptor.ofVoid(ValueLayout.JAVA_INT)); + original = ompMaxThreads(); + setOmpThreads(BASELINE); + } + + @After + public void restore() throws Throwable { + if (ompSetNumThreads != null) { + setOmpThreads(original); + } + } + + private int ompMaxThreads() throws Throwable { + return (int) ompGetMaxThreads.invokeExact(); + } + + private void setOmpThreads(int n) throws Throwable { + ompSetNumThreads.invokeExact(n); + } + + @Test + public void testCagraBuildRestoresOmpThreads() throws Throwable { + var params = + new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT) + .withNumWriterThreads(3) + .build(); + try (var resources = CheckedCuVSResources.create(); + var dataset = CuVSMatrix.ofArray(createFloatMatrix(2000, 32))) { + try (var index = + CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(params).build()) { + assertEquals(BASELINE, ompMaxThreads()); + } + } + } + + @Test + public void testCagraBuildWithoutParamsLeavesOmpThreads() throws Throwable { + try (var resources = CheckedCuVSResources.create(); + var dataset = CuVSMatrix.ofArray(createFloatMatrix(2000, 32))) { + try (var index = CagraIndex.newBuilder(resources).withDataset(dataset).build()) { + assertEquals(BASELINE, ompMaxThreads()); + } + } + } + + @Test + public void testCagraMergeRestoresOmpThreads() throws Throwable { + var params = + new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT) + .withNumWriterThreads(3) + .build(); + try (var resources = CheckedCuVSResources.create(); + var dataset1 = CuVSMatrix.ofArray(createFloatMatrix(2000, 32)).toDevice(resources); + var dataset2 = CuVSMatrix.ofArray(createFloatMatrix(2000, 32)).toDevice(resources); + var index1 = + CagraIndex.newBuilder(resources).withDataset(dataset1).withIndexParams(params).build(); + var index2 = + CagraIndex.newBuilder(resources) + .withDataset(dataset2) + .withIndexParams(params) + .build()) { + setOmpThreads(BASELINE); + try (var merged = CagraIndex.merge(new CagraIndex[] {index1, index2}, params)) { + assertEquals(BASELINE, ompMaxThreads()); + } + } + } + + @Test + public void testBruteForceBuildWithoutParams() throws Throwable { + try (var resources = CheckedCuVSResources.create(); + var dataset = CuVSMatrix.ofArray(createFloatMatrix(500, 16))) { + try (var index = BruteForceIndex.newBuilder(resources).withDataset(dataset).build()) { + assertEquals(BASELINE, ompMaxThreads()); + } + } + } +} diff --git a/java/panama-bindings/headers.h b/java/panama-bindings/headers.h index 04ebcce742..9ee3ca257e 100644 --- a/java/panama-bindings/headers.h +++ b/java/panama-bindings/headers.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -17,3 +17,8 @@ void omp_set_num_threads(int n_writer_threads); * @brief function signature for getting omp threads */ int omp_get_num_threads(void); + +/** + * @brief function signature for getting the calling thread's OpenMP thread count + */ +int omp_get_max_threads(void);