From 589d784ce12d7a7c02e3498099ab694595eece61 Mon Sep 17 00:00:00 2001 From: Max Buckley Date: Mon, 12 Oct 2026 18:17:43 +0000 Subject: [PATCH 1/4] java: scope the OpenMP thread count to each build The CAGRA, BBQ and brute-force builds called omp_set_num_threads(n) and reset it to 1 afterwards, and forced 1 when no index parameters were given. OpenMP's nthreads-var is per calling thread (libgomp keeps it in the thread's implicit task), so concurrent builds on different threads never shared the setting; the cost was the count itself. Resetting to 1 left the calling thread single-threaded for any later OpenMP work, and a build without parameters ran its host phases on one thread. Add OmpThreadScope, which applies the requested count and restores the thread's previous value (omp_get_max_threads, newly bound) instead of 1, and leaves OpenMP alone when no count is given. Also fix two crashes on builds without index parameters: - CAGRA read IVF-PQ n_lists through the NULL ivf_pq_build_params of the default parameters and segfaulted the JVM; - brute force threw a NullPointerException from getNumWriterThreads(). Drop the host sync before cuvsCagraBuild: there is no work on the stream that the build (on the same stream) is not already ordered after. Adds OmpThreadsIT; all three tests fail on main. Co-Authored-By: Claude Opus 5.5 --- .../cuvs/internal/BruteForceIndexImpl.java | 11 +-- .../nvidia/cuvs/internal/CagraIndexImpl.java | 38 ++++----- .../cuvs/internal/common/OmpThreadScope.java | 45 +++++++++++ .../java/com/nvidia/cuvs/OmpThreadsIT.java | 81 +++++++++++++++++++ java/panama-bindings/headers.h | 5 ++ 5 files changed, 156 insertions(+), 24 deletions(-) create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/OmpThreadScope.java create mode 100644 java/cuvs-java/src/test/java/com/nvidia/cuvs/OmpThreadsIT.java 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..9261fbe643 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 @@ -26,6 +26,7 @@ import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.SearchResults; import com.nvidia.cuvs.internal.common.CloseableRMMAllocation; +import com.nvidia.cuvs.internal.common.OmpThreadScope; import com.nvidia.cuvs.internal.panama.cuvsFilter; import java.io.InputStream; import java.io.OutputStream; @@ -119,12 +120,14 @@ private IndexReference build( MemorySegment datasetMemSegment = dataset.memorySegment(); - omp_set_num_threads(bruteForceIndexParams.getNumWriterThreads()); - long datasetBytes = C_FLOAT_BYTE_SIZE * rows * cols; var index = createBruteForceIndex(); - try (var resourcesAccessor = resources.access()) { + // Index parameters are optional; without them, leave the OpenMP setting alone. + int numWriterThreads = + bruteForceIndexParams != null ? bruteForceIndexParams.getNumWriterThreads() : 0; + try (var ompThreads = OmpThreadScope.of(numWriterThreads); + var resourcesAccessor = resources.access()) { long cuvsResources = resourcesAccessor.handle(); try (var closeableDataMemorySegmentP = allocateRMMSegment(cuvsResources, datasetBytes)) { MemorySegment datasetMemorySegmentP = closeableDataMemorySegmentP.handle(); @@ -152,8 +155,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 491c5213fc..680d30b956 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; @@ -182,19 +183,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( @@ -203,9 +209,6 @@ private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixInterna var cuvsRes = resourcesAccessor.handle(); - // TODO: do we need a stream sync here? - checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); - datasetView = isCagraPaddedLayout(dataset) ? makePaddedViewHandle(cuvsRes, dataset, localArena) @@ -228,7 +231,6 @@ private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixInterna if (!success) { quietly(() -> checkCuVSError(cuvsCagraIndexDestroy(index), "cuvsCagraIndexDestroy")); } - quietly(() -> omp_set_num_threads(1)); } } } @@ -249,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( @@ -331,7 +332,6 @@ private IndexReference buildBbq( quietly( () -> checkCuVSError(cuvsBbqQuantizerDestroy(quantizer), "cuvsBbqQuantizerDestroy")); } - quietly(() -> omp_set_num_threads(1)); } } } 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..b89493f4a3 --- /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. + * 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..63c02e2257 --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/OmpThreadsIT.java @@ -0,0 +1,81 @@ +/* + * 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.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 { + + private MethodHandle ompGetMaxThreads; + + @Before + public void setup() { + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); + initializeRandom(); + var symbol = + SymbolLookup.libraryLookup("libgomp.so.1", Arena.global()).find("omp_get_max_threads"); + assumeTrue("libgomp symbols not visible", symbol.isPresent()); + ompGetMaxThreads = + Linker.nativeLinker() + .downcallHandle(symbol.get(), FunctionDescriptor.of(ValueLayout.JAVA_INT)); + } + + private int ompMaxThreads() throws Throwable { + return (int) ompGetMaxThreads.invokeExact(); + } + + @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))) { + int before = ompMaxThreads(); + try (var index = + CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(params).build()) { + assertEquals(before, ompMaxThreads()); + } + } + } + + @Test + public void testCagraBuildWithoutParamsLeavesOmpThreads() throws Throwable { + try (var resources = CheckedCuVSResources.create(); + var dataset = CuVSMatrix.ofArray(createFloatMatrix(2000, 32))) { + int before = ompMaxThreads(); + try (var index = CagraIndex.newBuilder(resources).withDataset(dataset).build()) { + assertEquals(before, ompMaxThreads()); + } + } + } + + @Test + public void testBruteForceBuildWithoutParams() throws Throwable { + try (var resources = CheckedCuVSResources.create(); + var dataset = CuVSMatrix.ofArray(createFloatMatrix(500, 16))) { + int before = ompMaxThreads(); + try (var index = BruteForceIndex.newBuilder(resources).withDataset(dataset).build()) { + assertEquals(before, ompMaxThreads()); + } + } + } +} diff --git a/java/panama-bindings/headers.h b/java/panama-bindings/headers.h index 04ebcce742..c2fba57eda 100644 --- a/java/panama-bindings/headers.h +++ b/java/panama-bindings/headers.h @@ -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); From 06c61253ee559d6caf4fd5f19266929d8be2e8bb Mon Sep 17 00:00:00 2001 From: Max Buckley Date: Thu, 22 Oct 2026 17:40:52 +0200 Subject: [PATCH 2/4] java: address review feedback on the OpenMP thread scope - Scope the OpenMP thread count in CagraIndexImpl.merge() as well, using the merge params' numWriterThreads. - Drop OpenMP handling from the brute-force build; nothing in that path uses OpenMP. - Restore the stream sync before cuvsCagraBuild; it is unrelated to this change and will be proposed separately. - OmpThreadsIT: fail instead of skipping when omp_get_max_threads cannot be resolved, and cover merge. - Run pre-commit (copyright headers, spotless). Co-Authored-By: Claude Opus 5.5 --- .../cuvs/internal/BruteForceIndexImpl.java | 7 +--- .../nvidia/cuvs/internal/CagraIndexImpl.java | 5 +++ .../cuvs/internal/common/OmpThreadScope.java | 2 +- .../java/com/nvidia/cuvs/OmpThreadsIT.java | 32 ++++++++++++++++--- java/panama-bindings/headers.h | 2 +- 5 files changed, 36 insertions(+), 12 deletions(-) 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 9261fbe643..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 @@ -26,7 +26,6 @@ import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.SearchResults; import com.nvidia.cuvs.internal.common.CloseableRMMAllocation; -import com.nvidia.cuvs.internal.common.OmpThreadScope; import com.nvidia.cuvs.internal.panama.cuvsFilter; import java.io.InputStream; import java.io.OutputStream; @@ -123,11 +122,7 @@ private IndexReference build( long datasetBytes = C_FLOAT_BYTE_SIZE * rows * cols; var index = createBruteForceIndex(); - // Index parameters are optional; without them, leave the OpenMP setting alone. - int numWriterThreads = - bruteForceIndexParams != null ? bruteForceIndexParams.getNumWriterThreads() : 0; - try (var ompThreads = OmpThreadScope.of(numWriterThreads); - var resourcesAccessor = resources.access()) { + try (var resourcesAccessor = resources.access()) { long cuvsResources = resourcesAccessor.handle(); try (var closeableDataMemorySegmentP = allocateRMMSegment(cuvsResources, datasetBytes)) { MemorySegment datasetMemorySegmentP = closeableDataMemorySegmentP.handle(); 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 680d30b956..2ec75e3e86 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 @@ -209,6 +209,9 @@ private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixInterna var cuvsRes = resourcesAccessor.handle(); + // TODO: do we need a stream sync here? + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); + datasetView = isCagraPaddedLayout(dataset) ? makePaddedViewHandle(cuvsRes, dataset, localArena) @@ -1238,7 +1241,9 @@ public static CagraIndex merge( var mergedIndex = createCagraIndex(); CagraIndexImpl merged = null; + int numWriterThreads = mergeParams != null ? mergeParams.getNumWriterThreads() : 0; try (var nativeMergeParams = segmentFromIndexParams(mergeParams); + var ompThreads = OmpThreadScope.of(numWriterThreads); var resourcesAccessor = resources.access()) { var cuvsRes = resourcesAccessor.handle(); 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 index b89493f4a3..7b0422d5ef 100644 --- 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 @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.internal.common; 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 index 63c02e2257..cf89a0eab6 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/OmpThreadsIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/OmpThreadsIT.java @@ -29,11 +29,11 @@ public void setup() { assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); var symbol = - SymbolLookup.libraryLookup("libgomp.so.1", Arena.global()).find("omp_get_max_threads"); - assumeTrue("libgomp symbols not visible", symbol.isPresent()); + SymbolLookup.libraryLookup("libgomp.so.1", Arena.global()) + .find("omp_get_max_threads") + .orElseThrow(); ompGetMaxThreads = - Linker.nativeLinker() - .downcallHandle(symbol.get(), FunctionDescriptor.of(ValueLayout.JAVA_INT)); + Linker.nativeLinker().downcallHandle(symbol, FunctionDescriptor.of(ValueLayout.JAVA_INT)); } private int ompMaxThreads() throws Throwable { @@ -68,6 +68,30 @@ public void testCagraBuildWithoutParamsLeavesOmpThreads() throws Throwable { } } + @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()) { + int before = ompMaxThreads(); + try (var merged = CagraIndex.merge(new CagraIndex[] {index1, index2}, params)) { + assertEquals(before, ompMaxThreads()); + } + } + } + @Test public void testBruteForceBuildWithoutParams() throws Throwable { try (var resources = CheckedCuVSResources.create(); diff --git a/java/panama-bindings/headers.h b/java/panama-bindings/headers.h index c2fba57eda..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 */ From 6a3982c4fd1a86bb434a7d54938f7dc241bfd25b Mon Sep 17 00:00:00 2001 From: Max Buckley Date: Thu, 22 Oct 2026 19:26:30 +0200 Subject: [PATCH 3/4] java: pin a known OpenMP baseline in OmpThreadsIT Compare against a fixed thread count (5) instead of whatever the test thread happened to have, so a regression that resets the count to 1 or leaves the requested count in place cannot pass. The merge test re-pins the baseline after building its inputs so it covers merge() alone. Co-Authored-By: Claude Opus 5.5 --- .../java/com/nvidia/cuvs/OmpThreadsIT.java | 48 +++++++++++++------ 1 file changed, 34 insertions(+), 14 deletions(-) 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 index cf89a0eab6..eb855a27b5 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/OmpThreadsIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/OmpThreadsIT.java @@ -14,6 +14,7 @@ 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; @@ -22,24 +23,46 @@ @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() { + public void setup() throws Throwable { assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); - var symbol = - SymbolLookup.libraryLookup("libgomp.so.1", Arena.global()) - .find("omp_get_max_threads") - .orElseThrow(); + var libgomp = SymbolLookup.libraryLookup("libgomp.so.1", Arena.global()); + var linker = Linker.nativeLinker(); ompGetMaxThreads = - Linker.nativeLinker().downcallHandle(symbol, FunctionDescriptor.of(ValueLayout.JAVA_INT)); + 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 = @@ -49,10 +72,9 @@ public void testCagraBuildRestoresOmpThreads() throws Throwable { .build(); try (var resources = CheckedCuVSResources.create(); var dataset = CuVSMatrix.ofArray(createFloatMatrix(2000, 32))) { - int before = ompMaxThreads(); try (var index = CagraIndex.newBuilder(resources).withDataset(dataset).withIndexParams(params).build()) { - assertEquals(before, ompMaxThreads()); + assertEquals(BASELINE, ompMaxThreads()); } } } @@ -61,9 +83,8 @@ public void testCagraBuildRestoresOmpThreads() throws Throwable { public void testCagraBuildWithoutParamsLeavesOmpThreads() throws Throwable { try (var resources = CheckedCuVSResources.create(); var dataset = CuVSMatrix.ofArray(createFloatMatrix(2000, 32))) { - int before = ompMaxThreads(); try (var index = CagraIndex.newBuilder(resources).withDataset(dataset).build()) { - assertEquals(before, ompMaxThreads()); + assertEquals(BASELINE, ompMaxThreads()); } } } @@ -85,9 +106,9 @@ public void testCagraMergeRestoresOmpThreads() throws Throwable { .withDataset(dataset2) .withIndexParams(params) .build()) { - int before = ompMaxThreads(); + setOmpThreads(BASELINE); try (var merged = CagraIndex.merge(new CagraIndex[] {index1, index2}, params)) { - assertEquals(before, ompMaxThreads()); + assertEquals(BASELINE, ompMaxThreads()); } } } @@ -96,9 +117,8 @@ public void testCagraMergeRestoresOmpThreads() throws Throwable { public void testBruteForceBuildWithoutParams() throws Throwable { try (var resources = CheckedCuVSResources.create(); var dataset = CuVSMatrix.ofArray(createFloatMatrix(500, 16))) { - int before = ompMaxThreads(); try (var index = BruteForceIndex.newBuilder(resources).withDataset(dataset).build()) { - assertEquals(before, ompMaxThreads()); + assertEquals(BASELINE, ompMaxThreads()); } } } From cd9be4e63a6f9feb45e26a6eee9e97b7f4b86402 Mon Sep 17 00:00:00 2001 From: Max Buckley Date: Sat, 24 Oct 2026 14:41:16 +0200 Subject: [PATCH 4/4] Java: open the merge OMP scope inside resources.access(); document numWriterThreads <= 0 Close the OpenMP thread scope before releasing the resources accessor, as build already does, so the previous count is restored on the thread that set it. Document that a non-positive numWriterThreads leaves the OpenMP default in place. Co-Authored-By: Claude Opus 5.5 --- .../java-api-com-nvidia-cuvs-cagraindexparams.md | 14 +++++++++----- .../java/com/nvidia/cuvs/CagraIndexParams.java | 6 +++++- .../com/nvidia/cuvs/internal/CagraIndexImpl.java | 4 ++-- 3 files changed, 16 insertions(+), 8 deletions(-) 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/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 2ec75e3e86..d2b7b54373 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 @@ -1243,8 +1243,8 @@ public static CagraIndex merge( CagraIndexImpl merged = null; int numWriterThreads = mergeParams != null ? mergeParams.getNumWriterThreads() : 0; try (var nativeMergeParams = segmentFromIndexParams(mergeParams); - var ompThreads = OmpThreadScope.of(numWriterThreads); - 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