n2tiled-improved-scaling-on-multithread

docs/src/strategy.md

@@ -17,3 +17,48 @@ a `strategy` option to switch to a different option.
 For ``e^+e^-`` algorithms particle densities are low, so the only
 implementation for these algorithms is effectively of the same type as
 `N2Plain`.
+
+## Reusing N2Tiled storage
+
+The ordinary `jet_reconstruct` and `tiled_jet_reconstruct` interfaces return an
+independently owned `ClusterSequence`. This is the simplest interface to use,
+but it does carry a performance penalty, which the workspace described below
+overcomes.
+
+High-throughput applications that process many events can avoid most per-event
+temporary allocations with an `N2TiledWorkspace`:
+
+```julia
+workspace = N2TiledWorkspace()
+
+for event in events
+    with_n2tiled_reconstruction(
+        workspace,
+        event;
+        algorithm = JetAlgorithm.AntiKt,
+        R = 0.4,
+    ) do clusterseq
+        # This selection is independently owned and can outlive the callback.
+        jets = inclusive_jets(clusterseq; ptmin = 5.0)
+
+        # To retain the complete clustering sequence instead:
+        # retained_clusterseq = deepcopy(clusterseq)
+    end
+end
+```
+
+The callback receives a complete `ClusterSequence`, so history, constituent,
+and exclusive-jet queries remain available. Its jets and history borrow storage
+from the workspace and are overwritten the next time that workspace is used.
+You **must copy** values that will outlive the callback or use the ordinary
+owning interface.
+
+A workspace must not be shared concurrently or used reentrantly. Parallel
+applications should create one workspace for each concurrent worker and keep
+each workspace owned by that worker. `release_n2tiled_workspace_capacity!` can
+be used after an unusually large event or when a long-lived worker should
+release retained storage.
+
+The [multithreaded N2Tiled example](https://github.com/JuliaHEP/JetReconstruction.jl/blob/main/examples/n2tiled-multithreaded.jl)
+demonstrates dynamic event scheduling with one workspace owned by each
+long-lived worker task.

examples/n2tiled-multithreaded.jl

@@ -0,0 +1,106 @@
+#! /usr/bin/env julia
+
+"""
+Process independent events concurrently with one reusable N2Tiled workspace
+owned by each long-lived worker task.
+
+Run with, for example:
+
+```sh
+julia --threads=auto --project=examples examples/n2tiled-multithreaded.jl \
+    --maxevents=100 test/data/events.pp13TeV.hepmc3.zst
+```
+"""
+
+using ArgParse
+using JetReconstruction
+using LorentzVectorHEP
+
+"""
+    threaded_n2tiled_reconstruction(events; R=0.4, ptmin=5.0)
+
+Reconstruct `events` using dynamically scheduled worker tasks. Each worker owns
+one `N2TiledWorkspace`; only independently owned inclusive-jet selections are
+stored after the callback returns.
+"""
+function threaded_n2tiled_reconstruction(events;
+                                         R::Real = 0.4,
+                                         ptmin::Real = 5.0)
+    selected_jets = Vector{Vector{LorentzVector{Float64}}}(undef, length(events))
+    isempty(events) && return selected_jets
+
+    nworkers = min(Threads.nthreads(), length(events))
+
+    # A shared channel balances events dynamically. Workspaces belong to the
+    # tasks below rather than to thread IDs, so task migration is safe.
+    jobs = Channel{Int}(length(events))
+    for event_index in eachindex(events)
+        put!(jobs, event_index)
+    end
+    close(jobs)
+
+    @sync for _ in 1:nworkers
+        Threads.@spawn begin
+            workspace = N2TiledWorkspace()
+
+            for event_index in jobs
+                selected_jets[event_index] = with_n2tiled_reconstruction(workspace,
+                                                                         events[event_index];
+                                                                         algorithm = JetAlgorithm.AntiKt,
+                                                                         R = R,) do clusterseq
+                    # `inclusive_jets` returns an independently owned vector.
+                    # Use `deepcopy(clusterseq)` here instead when the complete
+                    # clustering sequence must outlive this callback.
+                    inclusive_jets(clusterseq; ptmin = ptmin)
+                end
+            end
+        end
+    end
+
+    return selected_jets
+end
+
+function parse_command_line(args)
+    settings = ArgParseSettings(autofix_names = true)
+    @add_arg_table! settings begin
+        "--maxevents", "-n"
+        help = "Maximum number of events to read; -1 reads all events."
+        arg_type = Int
+        default = -1
+
+        "--ptmin"
+        help = "Minimum transverse momentum for inclusive jets."
+        arg_type = Float64
+        default = 5.0
+
+        "--distance", "-R"
+        help = "Jet radius parameter."
+        arg_type = Float64
+        default = 0.4
+
+        "file"
+        help = "HepMC3 event file to read."
+        required = true
+    end
+
+    return parse_args(args, settings; as_symbols = true)
+end
+
+function main(args = ARGS)
+    options = parse_command_line(args)
+    events = read_final_state_particles(options[:file], PseudoJet;
+                                        maxevents = options[:maxevents])
+
+    selected_jets = threaded_n2tiled_reconstruction(events;
+                                                    R = options[:distance],
+                                                    ptmin = options[:ptmin])
+
+    println("Processed $(length(events)) events with $(Threads.nthreads()) Julia threads; " *
+            "selected $(sum(length, selected_jets)) jets.")
+
+    return nothing
+end
+
+if abspath(PROGRAM_FILE) == @__FILE__
+    main()
+end

src/ClusterSequence.jl

@@ -19,6 +19,22 @@ const NonexistentParent = -2
 "Cluster recombined with beam"
 const BeamJet = -1
 
+# Julia 1.11 introduced the `shrink` keyword for `sizehint!`. Disable shrinking
+# when that keyword is available. Older Julia versions have no public no-shrink
+# sizehint API, so leave their capacity untouched and let `resize!`/`push!` grow
+# it as needed.
+@inline function _sizehint_for_reuse!(buffer::Vector,
+                                      requested::Integer)
+    requested >= 0 ||
+        throw(ArgumentError("requested capacity must be non-negative, got $requested"))
+
+    @static if VERSION >= v"1.11"
+        sizehint!(buffer, requested; shrink = false)
+    end
+
+    return buffer
+end
+
 """
     struct HistoryElement
 
@@ -69,35 +85,53 @@ function HistoryElement(jetp_index)
 end
 
 """
-    initial_history(particles)
+    initial_history!(
+        history::Vector{HistoryElement},
+        particles,
+    )
 
-Create an initial history for the given particles.
+Reset and initialise reusable clustering-history storage for `particles`.
 
-# Arguments
-- `particles`: The initial vector of stable particles.
+The vector's logical length is reset to the number of initial particles while
+retaining any capacity already owned by the vector.
 
-# Returns
-- `history`: An array of `HistoryElement` objects.
-- `Qtot`: The total energy in the event.
+The returned history vector is borrowed storage owned by the caller.
 """
-function initial_history(particles)
-    # reserve sufficient space for everything
-    history = Vector{HistoryElement}(undef, length(particles))
-    sizehint!(history, 2 * length(particles))
+function initial_history!(history::Vector{HistoryElement},
+                          particles)
+    N = length(particles)
 
-    Qtot::Float64 = 0
+    # Establish exactly N active initial-history slots. All slots are
+    # overwritten below, and `resize!` retains any existing excess capacity.
+    resize!(history, N)
 
-    for i in eachindex(particles)
+    Qtot::Float64 = 0.0
+
+    @inbounds for i in eachindex(particles)
         history[i] = HistoryElement(i)
 
-        # get cross-referencing right from the Jets
-        # particles[i]._cluster_hist_index = i
-        @assert cluster_hist_index(particles[i])==i "Cluster history index should match jet's index in the input vector. Expected $(i), got $(cluster_hist_index(particles[i]))"
+        @assert cluster_hist_index(particles[i])==i ("Cluster history index should match jet's index in the input vector. "*
+                                                     "Expected $(i), got $(cluster_hist_index(particles[i]))")
 
-        # determine the total energy in the event
         Qtot += particles[i].E
     end
-    history, Qtot
+
+    return history, Qtot
+end
+
+"""
+    initial_history(particles)
+
+Create independently owned initial clustering-history storage.
+"""
+function initial_history(particles)
+    # This is newly owned storage, so requesting the complete sequence size
+    # cannot discard reusable capacity. Preserve the original eager allocation
+    # on Julia versions that do not support `sizehint!(...; shrink=false)`.
+    history = Vector{HistoryElement}(undef, length(particles))
+    sizehint!(history, 2 * length(particles))
+
+    return initial_history!(history, particles)
 end
 
 """
@@ -216,29 +250,20 @@ function add_step_to_history!(clusterseq::ClusterSequence, parent1, parent2, jet
 end
 
 """
-    inclusive_jets(clusterseq::ClusterSequence{U}, ::Type{T} = LorentzVectorCyl{Float64}; ptmin = 0.0) where {T, U}
+    inclusive_jets(
+        clusterseq::ClusterSequence{U},
+        ::Type{T}=LorentzVector{Float64};
+        ptmin=0.0,
+    )
 
-Return all inclusive jets of a ClusterSequence with pt > ptmin.
+Return all inclusive jets of a `ClusterSequence` with transverse momentum
+greater than or equal to `ptmin`.
 
-# Arguments
-- `clusterseq::ClusterSequence`: The `ClusterSequence` object containing the
-  clustering history and jets.
-- `::Type{T} = LorentzVectorCyl{Float64}`: The return type used for the selected jets.
-- `ptmin::Float64 = 0.0`: The minimum transverse momentum (pt) threshold for the
-  inclusive jets.
-
-# Returns
-An array of `T` objects representing the inclusive jets.
-
-# Description
-This function computes the inclusive jets from a given `ClusterSequence` object.
-It iterates over the clustering history and checks the transverse momentum of
-each parent jet. If the transverse momentum is greater than or equal to `ptmin`,
-the jet is added to the array of inclusive jets.
+Valid return types are `LorentzVector`, `LorentzVectorCyl`, or the jet type of
+the input `clusterseq` (`U`, either `PseudoJet` or `EEJet` depending on the
+algorithm).
 
-Valid return types are `LorentzVector` `LorentzVectorCyl` or the jet type of the
-input `clusterseq` (`U` - either `PseudoJet` or `EEJet` depending which
-algorithm was used).
+The returned vector is independently owned.
 
 # Example
 ```julia

src/JetReconstruction.jl

@@ -76,7 +76,8 @@ export plain_jet_reconstruct
 include("TiledAlgoUtils.jl")
 # Algorithmic part, tiled reconstruction strategy with linked list jet objects
 include("TiledAlgoLL.jl")
-export tiled_jet_reconstruct
+export N2TiledWorkspace, release_n2tiled_workspace_capacity!, tiled_jet_reconstruct,
+       with_n2tiled_reconstruction
 
 ## E+E- algorithms
 include("EEAlgorithm.jl")

src/JetUtils.jl

@@ -138,3 +138,68 @@ function kt_scale(jet1::T, jet2::T) where {T <: FourMomentum}
     pt2 = JetReconstruction.pt(jet2)
     return min(pt1, pt2) * deltar(jet1, jet2)
 end
+
+"""
+    construct_reco_jets(particles, ::Type{J}, preprocess) where {J <: FourMomentum}
+
+Create independently owned, reconstruction-ready jets of type `J` from
+`particles`.
+
+When `preprocess` is `nothing`, inputs that already have type `J` are copied
+directly and other inputs are converted to `J`. Otherwise `preprocess` is called
+for every input particle.
+"""
+function construct_reco_jets(particles::AbstractVector{P},
+                             ::Type{J},
+                             preprocess) where {P, J <: FourMomentum}
+    TargetNumericalType = eltype(particles[1])
+    if TargetNumericalType <: Real
+        TargetJetType = concretize_return_type(J, TargetNumericalType)
+    else
+        TargetJetType = typeof(J(particles[1]))
+    end
+
+    recombination_particles = Vector{TargetJetType}()
+    sizehint!(recombination_particles, 2 * length(particles))
+
+    return construct_reco_jets!(recombination_particles,
+                                particles,
+                                preprocess)
+end
+
+"""
+    construct_reco_jets!(recombination_particles, particles, preprocess)
+
+Reset and fill reusable reconstruction-jet storage. The destination must not
+alias the input collection.
+"""
+function construct_reco_jets!(recombination_particles::Vector{J},
+                              particles::AbstractVector{P},
+                              preprocess) where {P, J <: FourMomentum}
+    Base.mightalias(recombination_particles, particles) &&
+        throw(ArgumentError("reusable jet storage must not alias the input particles"))
+
+    N = length(particles)
+    empty!(recombination_particles)
+    _sizehint_for_reuse!(recombination_particles, 2 * N)
+
+    if isnothing(preprocess)
+        if P === J
+            append!(recombination_particles, particles)
+        else
+            for (i, particle) in enumerate(particles)
+                push!(recombination_particles,
+                      J(particle; cluster_hist_index = i))
+            end
+        end
+    else
+        for (i, particle) in enumerate(particles)
+            push!(recombination_particles,
+                  preprocess(particle,
+                             J;
+                             cluster_hist_index = i))
+        end
+    end
+
+    return recombination_particles
+end