From 8be2e61ef8df33a31fef694c00ff6806a193710a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Mon, 20 Oct 2025 16:14:02 +0200
Subject: [PATCH] Add interface for managing threads

---
 docs/src/api.md  |  7 +++++++
 src/HYPRE.jl     | 44 +++++++++++++++++++++++++++++++++++++++++---
 src/LibHYPRE.jl  |  4 ++++
 test/runtests.jl |  8 ++++++++
 4 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index ca6b8fb..b206d07 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -6,6 +6,13 @@
 HYPRE.Init
 ```
 
+## Threads
+
+```@docs
+HYPRE.set_nthreads()
+HYPRE.nthreads()
+```
+
 ## Matrix/vector creation
 
 ```@docs
diff --git a/src/HYPRE.jl b/src/HYPRE.jl
index 002adf1..ae2b53c 100644
--- a/src/HYPRE.jl
+++ b/src/HYPRE.jl
@@ -10,27 +10,37 @@ export HYPREMatrix, HYPREVector
 # Clang.jl auto-generated bindings and some manual methods
 include("LibHYPRE.jl")
 using .LibHYPRE
-using .LibHYPRE: @check
+using .LibHYPRE: @check, HYPRE_SetNumThreads, HYPRE_NumThreads
 
 # Internal namespace to hide utility functions
 include("Internals.jl")
 
 
 """
-    Init(; finalize_atexit=true)
+    Init(; finalize_atexit=true, nthreads = 1)
 
 Wrapper around `HYPRE_Initialize`. If `finalize_atexit` is `true` a Julia exit hook is added,
 which calls `HYPRE_Finalize`. This method will also call `MPI.Init` unless MPI is already
 initialized.
 
+The optional argument `nthreads` can be used to set the number of OpenMP threads
+HYPRE should use. The default is `1`, meaning no multithreading. See
+[`set_nthreads`](@ref) for more details. Setting `threads` to `0` or a negative
+value means that the number of threads will be controlled by hypre internally.
+This will typically be equal to the maximum number of threads, or the value in
+the ENV variable `OMP_NUM_THREADS`.
+
 **Note**: This function *must* be called before using HYPRE functions.
 """
-function Init(; finalize_atexit = true)
+function Init(; nthreads = 1, finalize_atexit = true)
     if !(MPI.Initialized())
         MPI.Init()
     end
     # TODO: Check if already initialized?
     HYPRE_Initialize()
+    if nthreads > 0
+        set_nthreads(nthreads)
+    end
     if finalize_atexit
         # TODO: MPI only calls the finalizer if not exiting due to a Julia exeption. Does
         #       the same reasoning apply here?
@@ -44,6 +54,34 @@ function Init(; finalize_atexit = true)
     return nothing
 end
 
+"""
+    set_nthreads(1)  # Single OpenMP thread
+    set_nthreads(10) # Use up to 10 OpenMP threads
+
+Set the number of OpenMP threads HYPRE should use internally on current process.
+The default is `1`, meaning no multithreading. Setting it to a zero or a
+negative value means that the number of threads will be unchanged.
+
+**Note***: The number of threads can improve execution speed, but a large number
+of threads can be detrimental to actual solver performance for some solvers
+(e.g. parallel Gauss-Seidel smoothers).
+"""
+function set_nthreads(nt::Integer)
+    if nt > 0
+        nt = min(nt, Sys.CPU_THREADS)
+        HYPRE_SetNumThreads(nt)
+    end
+    return nthreads()
+end
+
+"""
+    n = threads()
+
+Get the current number of OpenMP threads HYPRE is set to use on current process.
+"""
+function nthreads()
+    return HYPRE_NumThreads()
+end
 
 ###############
 # HYPREMatrix #
diff --git a/src/LibHYPRE.jl b/src/LibHYPRE.jl
index 4fef1f2..4247198 100644
--- a/src/LibHYPRE.jl
+++ b/src/LibHYPRE.jl
@@ -19,6 +19,10 @@ function HYPRE_SetNumThreads(nt::HYPRE_Int)
     return @ccall libHYPRE.hypre_SetNumThreads(nt::HYPRE_Int)::Ptr{Cvoid}
 end
 
+function HYPRE_SetNumThreads(nt::Integer)
+    return HYPRE_SetNumThreads(HYPRE_Int(nt))
+end
+
 # Add manual methods for some ::Function signatures where the library wants function
 # pointers. Instead of creating function pointers to the Julia wrappers we can just look
 # up the pointer in the library and pass that.
diff --git a/test/runtests.jl b/test/runtests.jl
index 726d186..c5ff487 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -29,6 +29,14 @@ end
     @test H.parmatrix != HYPRE_ParCSRMatrix(C_NULL)
 end
 
+@testset "Threads" begin
+    @test HYPRE.set_nthreads(1) == 1
+    @test HYPRE.set_nthreads(2) == 2
+    @test HYPRE.nthreads() == 2
+    @test HYPRE.set_nthreads(0) == 2
+    @test HYPRE.set_nthreads(1_000_000) == Sys.CPU_THREADS
+end
+
 @testset "HYPREMatrix(::SparseMatrixCS(C|R))" begin
     ilower, iupper = 4, 6
     CSC = convert(