Repository: ajwheeler/Korg.jl
Pull request #413: Threading
Author: login: ajwheeler
Created: 2025-04-16T19:41:45Z | Merged: 2025-04-22T15:40:41Z

This PR introduces thread parallelism to speed up parts of the calculation.  In the past, my philosophy was that Korg should be entirely serial, because you can always parallelize over stars/spectra.  However, it's still nice to have a single synthesis go faster on your laptop, and multithreading makes it easier to saturate the cores on a node without bumping into the memory ceiling.

[See here for how to set the number of julia threads](https://docs.julialang.org/en/v1/manual/multi-threading/) (it's `-t` or `$JULIA_NUM_THEADS`).  If you do nothing, you will continue to get the serial behavior which should have the same performance as before.


Diff:
diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml
index 16df6ec7..35a7f90a 100644
--- a/.github/workflows/Test.yml
+++ b/.github/workflows/Test.yml
@@ -30,6 +30,8 @@ jobs:
           arch: ${{ matrix.arch }}
       - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
+      - name: Set Julia threads
+        run: echo "JULIA_NUM_THREADS=4" >> $GITHUB_ENV
       - uses: julia-actions/julia-runtest@v1
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v5
diff --git a/.github/workflows/TestNightly.yml b/.github/workflows/TestNightly.yml
index 2a3c1286..d39ed41c 100644
--- a/.github/workflows/TestNightly.yml
+++ b/.github/workflows/TestNightly.yml
@@ -30,4 +30,6 @@ jobs:
           arch: ${{ matrix.arch }}
       - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
-      - uses: julia-actions/julia-runtest@v1
\ No newline at end of file
+      - name: Set Julia threads
+        run: echo "JULIA_NUM_THREADS=4" >> $GITHUB_ENV
+      - uses: julia-actions/julia-runtest@v1
diff --git a/docs/src/FAQ.md b/docs/src/FAQ.md
index 3e0e091e..9da89905 100644
--- a/docs/src/FAQ.md
+++ b/docs/src/FAQ.md
@@ -1,6 +1,16 @@
 ### What's the difference between `Korg.species` and `Korg.Species`?
-The short answer is that they do the same thing, but the lower-case version can only be applied to string literals.
-Both of these are ways to constructing objects of type [`Korg.Species`](@ref), which is used to represent atom and molecules with specific charges, e.g. C II or neutral FeH.  
+Short answer: you can always use `Korg.Species` like so: `Korg.Species("C II").
+Longer answer: they do the same thing, but the lower-case version can only be applied to string literals.
+Both of these are ways to constructing objects of type [`Korg.Species`](@ref), which is used to represent atom and molecules with specific charges, e.g. C II or neutral FeH.
 `Korg.Species` is the constructor, which acts like a normal function.  `Korg.species` is a [non-standard string literal](https://docs.julialang.org/en/v1/manual/metaprogramming/#meta-non-standard-string-literals), which is used without parentheses, like this: `Korg.species"Mg I"`.
 It can only be applied to string literals (things that you actually type out between quotes, not variables containing strings.)
 The string macro does the work of constructing the species object at compile time, which can make code much faster in specific circumstances.
+
+### Are Korg spectra bitwise reproducible?
+Korg spectra should be bitwise reproducible under these conditions:
+- Using a single thread (Julia's default configuration)
+- Running on the same machine
+- Using identical Korg and dependency versions
+- Providing identical input parameters
+
+When running with multiple threads, small numerical differences may occur due to floating-point operations being performed in different orders.
diff --git a/src/line_absorption.jl b/src/line_absorption.jl
index 8855f529..b1ad4a50 100644
--- a/src/line_absorption.jl
+++ b/src/line_absorption.jl
@@ -1,5 +1,6 @@
 using SpecialFunctions: gamma
 using ProgressMeter: @showprogress
+using Base.Iterators: partition
 
 """
     line_absorption!(α, linelist, λs, temp, nₑ, n_densities, partition_fns, ξ
@@ -22,18 +23,16 @@ other arguments:
 # Keyword Arguments
 
   - `cuttoff_threshold` (default: 3e-4): see `α_cntm`
-  - `verbose` (default: false): if true, show a progress bar.
+  - `tasks_per_thread` (default: 1): the number of tasks to run per Julia thread. This function
+    is multithreaded over the lines in `linelist`.
+  - `verbose` (deprecated): no longer used.
 """
 function line_absorption!(α, linelist, λs::Wavelengths, temps, nₑ, n_densities, partition_fns, ξ,
-                          α_cntm; cutoff_threshold=3e-4, verbose=false)
+                          α_cntm; cutoff_threshold=3e-4, verbose=false, tasks_per_thread=1)
     if length(linelist) == 0
         return zeros(length(λs))
     end
 
-    #lb and ub are the indices to the upper and lower wavelengths in the "window", i.e. the shortest
-    #and longest wavelengths which feel the effect of each line 
-    lb = 1
-    ub = 1
     β = @. 1 / (kboltz_eV * temps)
 
     # precompute number density / partition function for each species in the linelist
@@ -44,58 +43,74 @@ function line_absorption!(α, linelist, λs::Wavelengths, temps, nₑ, n_densiti
         @error "Atomic hydrogen should not be in the linelist. Korg has built-in hydrogen lines."
     end
 
-    # preallocate some arrays for the core loop. 
-    # Each element of the arrays corresponds to an atmospheric layer, same at the "temps" array and 
-    # the values in "number_densities"
-    Γ = Vector{eltype(α)}(undef, size(temps))
-    γ = Vector{eltype(α)}(undef, size(temps))
-    σ = Vector{eltype(α)}(undef, size(temps))
-    amplitude = Vector{eltype(α)}(undef, size(temps))
-    levels_factor = Vector{eltype(α)}(undef, size(temps))
-    ρ_crit = Vector{eltype(α)}(undef, size(temps))
-    inverse_densities = Vector{eltype(α)}(undef, size(temps))
-    @showprogress desc="calculating line opacities" enabled=verbose for line in linelist
-        m = get_mass(line.species)
-
-        # doppler-broadening width, σ (NOT √[2]σ)
-        σ .= doppler_width.(line.wl, temps, m, ξ)
-
-        # sum up the damping parameters.  These are FWHM (γ is usually the Lorentz HWHM) values in 
-        # angular, not cyclical frequency (ω, not ν).
-        Γ .= line.gamma_rad
-        if !ismolecule(line.species)
-            @. Γ += nₑ * scaled_stark.(line.gamma_stark, temps)
-            Γ .+= n_densities[species"H_I"] .* scaled_vdW.(Ref(line.vdW), m, temps)
+    n_chunks = tasks_per_thread * Threads.nthreads()
+    chunk_size = max(1, length(linelist) ÷ n_chunks + (length(linelist) % n_chunks > 0))
+    linelist_chunks = partition(linelist, chunk_size)
+    tasks = map(linelist_chunks) do linelist_chunk
+        # Each chunk of your data gets its own spawned task that does its own local, sequential work
+        # and then returns the result
+        Threads.@spawn begin
+            α_task = zeros(eltype(α), size(α))
+
+            # preallocate some arrays for the core loop.
+            # Each element of the arrays corresponds to an atmospheric layer, same at the "temps" array and
+            # the values in "number_densities"
+            Γ = Vector{eltype(α)}(undef, size(temps))
+            γ = Vector{eltype(α)}(undef, size(temps))
+            σ = Vector{eltype(α)}(undef, size(temps))
+            amplitude = Vector{eltype(α)}(undef, size(temps))
+            levels_factor = Vector{eltype(α)}(undef, size(temps))
+            ρ_crit = Vector{eltype(α)}(undef, size(temps))
+            inverse_densities = Vector{eltype(α)}(undef, size(temps))
+
+            for line in linelist_chunk
+                m = get_mass(line.species)
+
+                # doppler-broadening width, σ (NOT √[2]σ)
+                σ .= doppler_width.(line.wl, temps, m, ξ)
+
+                # sum up the damping parameters.  These are FWHM (γ is usually the Lorentz HWHM) values in
+                # angular, not cyclical frequency (ω, not ν).
+                Γ .= line.gamma_rad
+                if !ismolecule(line.species)
+                    @. Γ += nₑ * scaled_stark.(line.gamma_stark, temps)
+                    Γ .+= n_densities[species"H_I"] .* scaled_vdW.(Ref(line.vdW), m, temps)
+                end
+                # calculate the lorentz broadening parameter in wavelength. Doing this involves an
+                # implicit aproximation that λ(ν) is linear over the line window.
+                # the factor of λ²/c is |dλ/dν|, the factor of 1/2π is for angular vs cyclical freqency,
+                # and the last factor of 1/2 is for FWHM vs HWHM
+                @. γ = Γ * line.wl^2 / (c_cgs * 4π)
+
+                E_upper = line.E_lower + c_cgs * hplanck_eV / line.wl
+                @. levels_factor = exp(-β * line.E_lower) - exp(-β * E_upper)
+
+                #total wl-integrated absorption coefficient
+                @. amplitude = 10.0^line.log_gf * sigma_line(line.wl) * levels_factor *
+                               n_div_Z[line.species]
+
+                ρ_crit .= (line.wl .|> α_cntm) .* cutoff_threshold ./ amplitude
+                inverse_densities .= inverse_gaussian_density.(ρ_crit, σ)
+                doppler_line_window = maximum(inverse_densities)
+                inverse_densities .= inverse_lorentz_density.(ρ_crit, γ)
+                lorentz_line_window = maximum(inverse_densities)
+                window_size = sqrt(lorentz_line_window^2 + doppler_line_window^2)
+                # at present, this line is allocating. Would be good to fix that.
+                lb = searchsortedfirst(λs, line.wl - window_size)
+                ub = searchsortedlast(λs, line.wl + window_size)
+                # not necessary, but is faster as of 8f979cc2c28f45cd7230d9ee31fbfb5a5164eb1d
+                if lb > ub
+                    continue
+                end
+
+                α_task[:, lb:ub] .+= line_profile.(line.wl, σ, γ, amplitude, view(λs, lb:ub)')
+            end
+            return α_task
         end
-        # calculate the lorentz broadening parameter in wavelength. Doing this involves an 
-        # implicit aproximation that λ(ν) is linear over the line window.
-        # the factor of λ²/c is |dλ/dν|, the factor of 1/2π is for angular vs cyclical freqency,
-        # and the last factor of 1/2 is for FWHM vs HWHM
-        @. γ = Γ * line.wl^2 / (c_cgs * 4π)
-
-        E_upper = line.E_lower + c_cgs * hplanck_eV / line.wl
-        @. levels_factor = exp(-β * line.E_lower) - exp(-β * E_upper)
-
-        #total wl-integrated absorption coefficient
-        @. amplitude = 10.0^line.log_gf * sigma_line(line.wl) * levels_factor *
-                       n_div_Z[line.species]
-
-        ρ_crit .= (line.wl .|> α_cntm) .* cutoff_threshold ./ amplitude
-        inverse_densities .= inverse_gaussian_density.(ρ_crit, σ)
-        doppler_line_window = maximum(inverse_densities)
-        inverse_densities .= inverse_lorentz_density.(ρ_crit, γ)
-        lorentz_line_window = maximum(inverse_densities)
-        window_size = sqrt(lorentz_line_window^2 + doppler_line_window^2)
-        # at present, this line is allocating. Would be good to fix that.
-        lb = searchsortedfirst(λs, line.wl - window_size)
-        ub = searchsortedlast(λs, line.wl + window_size)
-        # not necessary, but is faster as of 8f979cc2c28f45cd7230d9ee31fbfb5a5164eb1d
-        if lb > ub
-            continue
-        end
-
-        view(α, :, lb:ub) .+= line_profile.(line.wl, σ, γ, amplitude, view(λs, lb:ub)')
     end
+
+    α .+= sum(fetch.(tasks))
+    return nothing
 end
 
 """
diff --git a/test/runtests.jl b/test/runtests.jl
index 4a8d87a8..cef165eb 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,5 +1,7 @@
 using Korg, Test, Logging, HDF5, ForwardDiff, FiniteDiff, Aqua
 
+println("running test suite with  ", Threads.nthreads(), " threads.")
+
 @testset "Korg tests" verbose=true showtiming=true begin
 
     # tools for testing: assert_allclose and assert_allclose_grid


Reviews:
user: login: mabruzzo
body: Neat!

I made a quick suggestion for marginally better partitioning.

I would also strongly recommend that you add a note (maybe just a sentence to the docstring and perhaps to the general documentation) telling people that results may not be bitwise identical run-to-run when using multiple threads (due to the nonassociativity of floating-point addition). Obviously, this only becomes a problem if lines-widths are broad compared to the range of line-centers per chunk. Even if this seems totally obvious to you/me, I really think its worth mentioning. It can totally catch people by surprise
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-18T00:37:59Z
user: login: ajwheeler
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-18T13:54:32Z

Review comments:
diff_hunk: @@ -44,58 +43,73 @@ function line_absorption!(α, linelist, λs::Wavelengths, temps, nₑ, n_densiti
         @error "Atomic hydrogen should not be in the linelist. Korg has built-in hydrogen lines."
     end
 
-    # preallocate some arrays for the core loop. 
-    # Each element of the arrays corresponds to an atmospheric layer, same at the "temps" array and 
-    # the values in "number_densities"
-    Γ = Vector{eltype(α)}(undef, size(temps))
-    γ = Vector{eltype(α)}(undef, size(temps))
-    σ = Vector{eltype(α)}(undef, size(temps))
-    amplitude = Vector{eltype(α)}(undef, size(temps))
-    levels_factor = Vector{eltype(α)}(undef, size(temps))
-    ρ_crit = Vector{eltype(α)}(undef, size(temps))
-    inverse_densities = Vector{eltype(α)}(undef, size(temps))
-    @showprogress desc="calculating line opacities" enabled=verbose for line in linelist
-        m = get_mass(line.species)
-
-        # doppler-broadening width, σ (NOT √[2]σ)
-        σ .= doppler_width.(line.wl, temps, m, ξ)
-
-        # sum up the damping parameters.  These are FWHM (γ is usually the Lorentz HWHM) values in 
-        # angular, not cyclical frequency (ω, not ν).
-        Γ .= line.gamma_rad
-        if !ismolecule(line.species)
-            @. Γ += nₑ * scaled_stark.(line.gamma_stark, temps)
-            Γ .+= n_densities[species"H_I"] .* scaled_vdW.(Ref(line.vdW), m, temps)
+    chunk_size = max(1, length(linelist) ÷ (tasks_per_thread * Threads.nthreads()))
path: src/line_absorption.jl
user: login: mabruzzo
body: Let `n_chunks` denote `tasks_per_thread * Threads.nthreads()`.

As currently written, the code only splits `linelist` into `n_chunks` iff `length(linelist)` is an integer multiple of `n_chunks`. In the more common case where `length(linelist) % n_chunks > 0`, `linelist` will be split into `n_chunks+1`. In the limit where each chunk is the same amount of work, one of of your threads will have an extra task.

The following tweak ensures that you always split `linelist` into `n_chunks`, when `length(linelist) % n_chunks > 0`. (As you're probably aware, this is modification is cheap since most compilers know how to group nearby divisions and remainder operations into a single instruction).

```suggestion
    n_chunks = tasks_per_thread * Threads.nthreads()
    chunk_size = max(1, length(linelist) ÷ n_chunks + (length(linelist) % n_chunks > 0))
```
created_at: 2025-04-18T00:25:10Z
author_association: COLLABORATOR
diff_hunk: @@ -44,58 +43,73 @@ function line_absorption!(α, linelist, λs::Wavelengths, temps, nₑ, n_densiti
         @error "Atomic hydrogen should not be in the linelist. Korg has built-in hydrogen lines."
     end
 
-    # preallocate some arrays for the core loop. 
-    # Each element of the arrays corresponds to an atmospheric layer, same at the "temps" array and 
-    # the values in "number_densities"
-    Γ = Vector{eltype(α)}(undef, size(temps))
-    γ = Vector{eltype(α)}(undef, size(temps))
-    σ = Vector{eltype(α)}(undef, size(temps))
-    amplitude = Vector{eltype(α)}(undef, size(temps))
-    levels_factor = Vector{eltype(α)}(undef, size(temps))
-    ρ_crit = Vector{eltype(α)}(undef, size(temps))
-    inverse_densities = Vector{eltype(α)}(undef, size(temps))
-    @showprogress desc="calculating line opacities" enabled=verbose for line in linelist
-        m = get_mass(line.species)
-
-        # doppler-broadening width, σ (NOT √[2]σ)
-        σ .= doppler_width.(line.wl, temps, m, ξ)
-
-        # sum up the damping parameters.  These are FWHM (γ is usually the Lorentz HWHM) values in 
-        # angular, not cyclical frequency (ω, not ν).
-        Γ .= line.gamma_rad
-        if !ismolecule(line.species)
-            @. Γ += nₑ * scaled_stark.(line.gamma_stark, temps)
-            Γ .+= n_densities[species"H_I"] .* scaled_vdW.(Ref(line.vdW), m, temps)
+    chunk_size = max(1, length(linelist) ÷ (tasks_per_thread * Threads.nthreads()))
path: src/line_absorption.jl
user: login: ajwheeler
body: Great catch, thanks!
created_at: 2025-04-18T13:54:32Z
author_association: OWNER

PR comments:
user: login: ajwheeler
created_at: 2025-04-18T18:11:53Z
body: Added a note about reproducibility to the FAQ.  An explanation of multithreading is also in #416, which I will merge after releasing an update.
author_association: OWNER