toplevel: Parallelize report_package signature analysis (#785)

Parallelize the analysis of method signatures in `report_package` using
`Threads.@spawn`. Each `analyze_method_signature!` task gets its own
independent `AnalyzerState` with local `inf_cache` and `analysis_results`,
avoiding data races without complex cross-task cache sharing.

Design notes: https://publish.obsidian.md/jetls/work/JETLS/JET/JET+Make+%60report_package%60+parallelized

Key changes:
- Add `AtomicContainers` module providing thread-safe container types
  (`SWContainer`, `LWContainer`, `CASContainer`) for concurrent access
- Introduce `CASDict` for thread-safe analyzer cache (`JET_ANALYZER_CACHE`
  and `OPT_ANALYZER_CACHE`)
- Add `PackageAnalysisProgress` struct for tracking parallel analysis
  progress with atomic counters
- Parallelize `analyze_from_definitions!` in virtualprocess.jl

Benchamrks:

> `julia --startup-file-no --threads=4,2 -e 'using JET; report_package(JET; target_modules=(JET,), sourceinfo=:compact);'`

| Approach   | Time   |
|------------|--------|
| sequential | 52.07s |
| parallel   | 17.75s |

> ``julia --startup-file=no --threads=4,2 -e 'using JET; using Pkg; Pkg.activate(; temp=true); Pkg.add("CSV"); using CSV; report_package(CSV; target_modules=(CSV,), sourceinfo=:compact);'``

| Approach   | Time   |
|------------|--------|
| sequential | 44.23s |
| parallel   | 19.57s |

Author: aviatesk

.github/workflows/ci.yml

@@ -83,6 +83,26 @@ jobs:
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v5
 
+  multithreading-test:
+    name: JULIA_NUM_THREADS=4,2
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: "1"
+          arch: x64
+      - uses: julia-actions/cache@v2
+      - name: set preferences
+        working-directory: .
+        run: |
+          echo '[JET]
+          JET_DEV_MODE = true' >  LocalPreferences.toml
+      - uses: julia-actions/julia-buildpkg@latest
+      - uses: julia-actions/julia-runtest@latest
+        env:
+          JULIA_NUM_THREADS: "4,2"
+
   empty-loading-test:
     runs-on: ubuntu-latest
     steps:

CHANGELOG.md

@@ -53,6 +53,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 <!-- links end -->
 
 ## [Unreleased]
+### Changed
+- **Parallelized `report_package`**: Method signature analysis in `report_package`
+  is now parallelized using Julia's multithreading, providing significant
+  speedup on multi-core systems.
+
+  > With `--threads=4,2`
+  - Benchmark on `report_package(JET)`: 52.07s → 17.75s (~3x faster)
+  - Benchmark on `report_package(CSV)`: 44.23s → 19.57s (~2x faster)
+
 ### Internal
 - Refactored the project file to use the [`[workspace]`](https://pkgdocs.julialang.org/v1/toml-files/#The-%5Bworkspace%5D-section) for the docs/test environment of JET.
   This allows running e.g. `julia --project=./test test/runtests.jl` or

src/JETBase.jl

@@ -294,6 +294,34 @@ const JET_LOGGER_LEVELS_DESC = let
 end
 jet_logger_level(@nospecialize io::IO) = get(io, JET_LOGGER_LEVEL, DEFAULT_LOGGER_LEVEL)::Int
 
+# multithreading
+
+"""
+    CASDict{K,V}
+
+A thread-safe dictionary using Compare-And-Swap (CAS) operations for lock-free updates.
+Reads are always lock-free via atomic load. Writes use a CAS retry loop, making this
+suitable for lightweight, pure update functions that are safe to retry.
+
+Currently implements [`get!`](@ref) only.
+"""
+mutable struct CASDict{K,V}
+    @atomic data::Dict{K,V}
+    CASDict{K,V}() where {K,V} = new{K,V}(Dict{K,V}())
+end
+
+function Base.get!(f::Base.Callable, d::CASDict{K,V}, key::K) where {K,V}
+    old = @atomic :acquire d.data
+    val = get(old, key, nothing)
+    val !== nothing && return val::V
+    while true
+        new = copy(old)
+        val = get!(f, new, key)::V
+        old, success = @atomicreplace :acquire_release :monotonic d.data old => new
+        success && return val
+    end
+end
+
 # analysis core
 # =============
 
@@ -327,6 +355,20 @@ Prints a report of the top-level error `report` to the given `io`.
 """
 function print_report end
 
+mutable struct PackageAnalysisProgress
+    const reports::Vector{InferenceErrorReport}
+    const reports_lock::ReentrantLock
+    @atomic done::Int
+    @atomic analyzed::Int
+    @atomic cached::Int
+    const interval::Int
+    @atomic next_interval::Int
+    function PackageAnalysisProgress(n_sigs::Int)
+        interval = max(n_sigs ÷ 25, 1)
+        new(InferenceErrorReport[], ReentrantLock(), 0, 0, 0, interval, interval)
+    end
+end
+
 include("toplevel/virtualprocess.jl")
 
 # results
@@ -968,6 +1010,11 @@ struct SigAnalysisResult
     codeinst::CodeInstance
 end
 
+struct SigWorkItem
+    siginfos::Vector{Revise.SigInfo}
+    index::Int
+end
+
 """
     analyze_and_report_package!(analyzer::AbstractAnalyzer, package::Module; jetconfigs...) -> JETToplevelResult
 
@@ -991,7 +1038,6 @@ function analyze_and_report_package!(analyzer::AbstractAnalyzer, pkgmod::Module;
     end
 
     start = time()
-    counter, analyzed, cached = Ref(0), Ref(0), Ref(0)
     res = VirtualProcessResult(nothing)
     jetconfigs = set_if_missing(jetconfigs, :toplevel_logger, IOContext(stdout, JET_LOGGER_LEVEL => DEFAULT_LOGGER_LEVEL))
     config = ToplevelConfig(; jetconfigs...)
@@ -1001,66 +1047,89 @@ function analyze_and_report_package!(analyzer::AbstractAnalyzer, pkgmod::Module;
     newstate = AnalyzerState(AnalyzerState(analyzer); world=Base.get_world_counter())
     analyzer = AbstractAnalyzer(analyzer, newstate)
 
-    n_sigs = 0
-    for fi in pkgdata.fileinfos, (_, exsigs) in fi.modexsigs, (_, siginfos) in exsigs
-        isnothing(siginfos) && continue
-        n_sigs += length(siginfos)
-    end
+    workitems = SigWorkItem[]
     for fi in pkgdata.fileinfos, (_, exsigs) in fi.modexsigs, (_, siginfos) in exsigs
         isnothing(siginfos) && continue
         for (i, siginfo) in enumerate(siginfos)
-            toplevel_logger(config) do @nospecialize(io::IO)
-                clearline(io)
-            end
-            counter[] += 1
-            inf_world = CC.get_inference_world(analyzer)
+            push!(workitems, SigWorkItem(siginfos, i))
+        end
+    end
+
+    n_sigs = length(workitems)
+    progress = PackageAnalysisProgress(n_sigs)
+    inf_world = CC.get_inference_world(analyzer)
+
+    toplevel_logger(config) do @nospecialize(io::IO)
+        print(io, "Analyzing top-level definitions (progress: 0/$n_sigs | interval: $(progress.interval))")
+    end
+
+    tasks = map(workitems) do workitem
+        (; siginfos, index) = workitem
+        siginfo = siginfos[index]
+        Threads.@spawn :default try
             ext = Revise.get_extended_data(siginfo, :JET)
+            local reports::Vector{InferenceErrorReport}
             if ext !== nothing && ext.data isa SigAnalysisResult
                 prev_result = ext.data::SigAnalysisResult
                 if (CC.cache_owner(analyzer) === prev_result.codeinst.owner &&
                     prev_result.codeinst.max_world ≥ inf_world ≥ prev_result.codeinst.min_world)
-                    toplevel_logger(config) do @nospecialize(io::IO)
-                        (counter[] == n_sigs ? println : print)(io, "Skipped analysis for cached definition ($(counter[])/$n_sigs)")
-                    end
-                    cached[] += 1
+                    @atomic progress.cached += 1
                     reports = prev_result.reports
                     @goto gotreports
                 end
             end
+            # Create a new analyzer with fresh local caches (`inf_cache` and `analysis_results`)
+            # to avoid data races between concurrent signature analysis tasks
+            task_analyzer = AbstractAnalyzer(analyzer,
+                AnalyzerState(AnalyzerState(analyzer), #=refresh_local_cache=#true))
             match = Base._which(siginfo.sig;
-                method_table = CC.method_table(analyzer),
+                method_table = CC.method_table(task_analyzer),
                 world = inf_world,
                 raise = false)
             if match !== nothing
-                toplevel_logger(config; pre=clearline) do @nospecialize(io::IO)
-                    if jet_logger_level(io) ≥ JET_LOGGER_LEVEL_DEBUG
-                        print(io, "Analyzing top-level definition `")
-                        Base.show_tuple_as_call(io, Symbol(""), siginfo.sig)
-                        print(io, "` (progress: $(counter[])/$n_sigs)")
-                    else
-                        print(io, "Analyzing top-level definition (progress: $(counter[])/$n_sigs)")
-                    end
-                end
-                result = analyze_method_signature!(analyzer,
+                result = analyze_method_signature!(task_analyzer,
                     match.method, match.spec_types, match.sparams)
-                analyzed[] += 1
-                reports = get_reports(analyzer, result)
-                siginfos[i] = Revise.replace_extended_data(siginfo, :JET, SigAnalysisResult(reports, result.ci))
-                @label gotreports
-                append!(res.inference_error_reports, reports)
+                @atomic progress.analyzed += 1
+                reports = get_reports(task_analyzer, result)
+                siginfos[index] = Revise.replace_extended_data(siginfo, :JET, SigAnalysisResult(reports, result.ci))
             else
-                toplevel_logger(config) do @nospecialize(io::IO)
+                toplevel_logger(config; pre=println) do @nospecialize(io::IO)
                     print(io, "Couldn't find a single matching method for the signature `")
                     Base.show_tuple_as_call(io, Symbol(""), siginfo.sig)
-                    println(io, "` (progress: $(counter[])/$n_sigs)")
+                    println(io, "`")
+                end
+                reports = InferenceErrorReport[]
+            end
+            @label gotreports
+            isempty(reports) || @lock progress.reports_lock append!(progress.reports, reports)
+        catch err
+            @error "Error analyzing method signature" siginfo.sig
+            Base.showerror(stderr, err, catch_backtrace())
+        finally
+            done = (@atomic progress.done += 1)
+            current_next = @atomic progress.next_interval
+            if done >= current_next
+                @atomicreplace progress.next_interval current_next => current_next + progress.interval
+                toplevel_logger(config; pre=clearline) do @nospecialize(io::IO)
+                    print(io, "Analyzing top-level definitions (progress: $done/$n_sigs)")
                 end
             end
         end
     end
 
+    waitall(tasks)
+
+    append!(res.inference_error_reports, progress.reports)
+
+    toplevel_logger(config; pre=clearline) do @nospecialize(io::IO)
+        done = @atomic progress.done
+        print(io, "Analyzing top-level definitions (progress: $done/$n_sigs)")
+    end
     toplevel_logger(config; pre=println) do @nospecialize(io::IO)
         sec = round(time() - start; digits = 3)
-        println(io, "Analyzed all top-level definitions (all: $(counter[]) | analyzed: $(analyzed[]) | cached: $(cached[]) | took: $sec sec)")
+        analyzed = @atomic progress.analyzed
+        cached = @atomic progress.cached
+        println(io, "Analyzed all top-level definitions (all: $n_sigs | analyzed: $analyzed | cached: $cached | took: $sec sec)")
     end
 
     unique!(aggregation_policy(analyzer), res.inference_error_reports)

src/abstractinterpret/abstractanalyzer.jl

@@ -295,7 +295,7 @@ end
 # -----------
 # 2. `AbstractAnalyzer(analyzer::NewAnalyzer, state::AnalyzerState) -> NewAnalyzer`
 
-@noinline function AbstractAnalyzer(analyzer::AbstractAnalyzer, state::AnalyzerState)
+@noinline function AbstractAnalyzer(analyzer::AbstractAnalyzer, ::AnalyzerState)
     AnalyzerType = nameof(typeof(analyzer))
     error(lazy"""
     Missing `$AbstractAnalyzer` API:

src/analyzers/jetanalyzer.jl

@@ -123,7 +123,7 @@ JETInterface.typeinf_world(::BasicJETAnalyzer) = JET_TYPEINF_WORLD[]
 JETInterface.typeinf_world(::SoundJETAnalyzer) = JET_TYPEINF_WORLD[]
 JETInterface.typeinf_world(::TypoJETAnalyzer) = JET_TYPEINF_WORLD[]
 
-const JET_ANALYZER_CACHE = Dict{UInt, AnalysisToken}()
+const JET_ANALYZER_CACHE = CASDict{UInt,AnalysisToken}()
 
 JETAnalyzerConfig(analyzer::JETAnalyzer) = analyzer.config
 

src/analyzers/optanalyzer.jl

@@ -160,7 +160,7 @@ JETInterface.AnalysisToken(analyzer::OptAnalyzer) = analyzer.analysis_token
 JETInterface.typeinf_world(::OptAnalyzer) = JET_TYPEINF_WORLD[]
 JETInterface.vscode_diagnostics_order(::OptAnalyzer) = false
 
-const OPT_ANALYZER_CACHE = Dict{UInt, AnalysisToken}()
+const OPT_ANALYZER_CACHE = CASDict{UInt,AnalysisToken}()
 
 # overloads
 # =========

src/toplevel/virtualprocess.jl

@@ -734,42 +734,71 @@ gen_virtual_module(parent::Module = Main; name = VIRTUAL_MODULE_NAME) =
 # if code generation has failed given the entry method signature, the overload of
 # `InferenceState(..., ::AbstractAnalyzer)` will collect `GeneratorErrorReport`
 function analyze_from_definitions!(interp::ConcreteInterpreter, config::ToplevelConfig)
-    succeeded = Ref(0)
     start = time()
-    analyzer = ToplevelAbstractAnalyzer(interp, non_toplevel_concretized; refresh_local_cache = false)
     entrypoint = config.analyze_from_definitions
     res = InterpretationState(interp).res
     n_sigs = length(res.toplevel_signatures)
-    for i = 1:n_sigs
-        tt = res.toplevel_signatures[i]
-        match = Base._which(tt;
-            # NOTE use the latest world counter with `method_table(analyzer)` unwrapped,
-            # otherwise it may use a world counter when this method isn't defined yet
-            method_table = CC.method_table(analyzer),
-            world = CC.get_inference_world(analyzer),
-            raise = false)
-        if (match !== nothing &&
-            (!(entrypoint isa Symbol) || # implies `analyze_from_definitions===true`
-             match.method.name === entrypoint))
-            succeeded[] += 1
-            toplevel_logger(config; pre=clearline) do @nospecialize(io::IO)
-                print(io, "analyzing from top-level definitions ($(succeeded[])/$n_sigs)")
+    n_sigs == 0 && return nothing
+
+    progress = PackageAnalysisProgress(n_sigs)
+
+    toplevel_logger(config) do @nospecialize(io::IO)
+        print(io, "analyzing from top-level definitions (0/$n_sigs)")
+    end
+
+    tasks = map(1:n_sigs) do i
+        Threads.@spawn begin
+            tt = res.toplevel_signatures[i]
+            # Create a new analyzer with fresh local caches (`inf_cache` and `analysis_results`)
+            # to avoid data races between concurrent signature analysis tasks
+            analyzer = ToplevelAbstractAnalyzer(interp, non_toplevel_concretized;
+                refresh_local_cache = true)
+            match = Base._which(tt;
+                # NOTE use the latest world counter with `method_table(analyzer)` unwrapped,
+                # otherwise it may use a world counter when this method isn't defined yet
+                method_table = CC.method_table(analyzer),
+                world = CC.get_inference_world(analyzer),
+                raise = false)
+            if (match !== nothing &&
+                (!(entrypoint isa Symbol) || # implies `analyze_from_definitions===true`
+                 match.method.name === entrypoint))
+                @atomic progress.analyzed += 1
+                result = analyze_method_signature!(analyzer,
+                    match.method, match.spec_types, match.sparams)
+                reports = get_reports(analyzer, result)
+                isempty(reports) || @lock progress.reports_lock append!(progress.reports, reports)
+            else
+                # something went wrong
+                toplevel_logger(config; pre=clearline) do @nospecialize(io::IO)
+                    println(io, "couldn't find a single method matching the signature `", tt, "`")
+                end
             end
-            result = analyze_method_signature!(analyzer,
-                match.method, match.spec_types, match.sparams)
-            reports = get_reports(analyzer, result)
-            append!(res.inference_error_reports, reports)
-        else
-            # something went wrong
-            toplevel_logger(config; filter=≥(JET_LOGGER_LEVEL_DEBUG), pre=clearline) do @nospecialize(io::IO)
-                println(io, "couldn't find a single method matching the signature `", tt, "`")
+            done = (@atomic progress.done += 1)
+            current_next = @atomic progress.next_interval
+            if done >= current_next
+                @atomicreplace progress.next_interval current_next => current_next + progress.interval
+                toplevel_logger(config; pre=clearline) do @nospecialize(io::IO)
+                    analyzed = @atomic progress.analyzed
+                    print(io, "analyzing from top-level definitions ($analyzed/$n_sigs)")
+                end
             end
         end
     end
+
+    waitall(tasks)
+
+    append!(res.inference_error_reports, progress.reports)
+
+    toplevel_logger(config; pre=clearline) do @nospecialize(io::IO)
+        done = @atomic progress.done
+        print(io, "analyzing from top-level definitions ($done/$n_sigs)")
+    end
     toplevel_logger(config; pre=println) do @nospecialize(io::IO)
         sec = round(time() - start; digits = 3)
-        println(io, "analyzed $(succeeded[]) top-level definitions (took $sec sec)")
+        analyzed = @atomic progress.analyzed
+        println(io, "analyzed $analyzed top-level definitions (took $sec sec)")
     end
+
     return nothing
 end