Java: scope the OpenMP thread count to each build; fix builds without index parameters

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`_

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
+     * <p>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) {

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);
     }
   }
 

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<MemorySegment>(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

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()}.
+ *
+ * <p>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);
+    }
+  }
+}

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());
+      }
+    }
+  }
+}

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);