Repository navigation
Conversation
hakkelt
force-pushed
the
unified-threading
branch
from
August 19, 2026 15:01
0287558 to
a8a806a
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #44 +/- ##
==========================================
- Coverage 90.13% 90.07% -0.06%
==========================================
Files 51 52 +1
Lines 3678 4142 +464
==========================================
+ Hits 3315 3731 +416
- Misses 363 411 +48 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
hakkelt
marked this pull request as ready for review
August 24, 2026 19:44
hakkelt
force-pushed
the
unified-threading
branch
2 times, most recently
from
August 27, 2026 05:11
525f485 to
c9c9fda
Compare
Collaborator
Author
|
@lostella, this PR is also ready for review. Sorry, this is quite a large PR again, but the goal is straightforward: add multithreading support for all operators where possible via a unified API (and fix minor issues encountered along the way while implementing the feature). The proposed changes don't break the current API, and the default value of the parameter governing multithreading aims to maximize out-of-the-box performance. Based on the operator type and input size, it enables multithreading only when preliminary benchmarks indicate it is beneficial. The goal of adding the parameter named This is the suggested unified API: op1 = Variation(Float64, (64,64))
x, y = rand(64, 64), (4096, 2)
mul!(y, op1, x) # this runs multithreaded
op2 = Variation(Float64, (64,64); threaded=false)
x, y = rand(64, 64), (4096, 2)
mul!(y, op2, x) # this runs single-threaded
op3 = Variation(Float64, (10,10))
x, y = rand(10, 10), (100, 2)
mul!(y, op3, x) # this runs single-threaded because the input is too small to benefit from multithreading (2^10 elements is the threshold for Variation)A side-effect of this work is that the |
lostella
self-requested a review
August 29, 2026 05:17
lostella
reviewed
Sep 13, 2026
hakkelt
referenced
this pull request
in hakkelt/AbstractOperators.jl
Sep 19, 2026
- DFT.jl: use _with_fftw_threads (try/finally) for the c2c plan construction instead of raw set_num_threads/get_num_threads, matching DCT.jl/RDFT.jl/IRDFT.jl. - DCAT.jl: forward threaded = is_block_threaded(H) in permute and remove_displacement reconstruction, matching VCAT.jl precedent. - HCAT.jl: forward threaded = is_block_threaded(L) in remove_slicing (permute/remove_displacement already did this). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
hakkelt
force-pushed
the
unified-threading
branch
from
September 19, 2026 13:00
44543d6 to
4c0ab10
Compare
Benchmark Results (Julia v1.13.0)Run once per thread count ( 1 thread🚀 8 benchmarks improved in time · 🐢 1 time regression detected · 🚀 7 benchmarks use less memory Time benchmarks
Memory benchmarks
2 threads🚀 25 benchmarks improved in time · 🚀 6 benchmarks use less memory · 🐢 3 memory regressions detected Time benchmarks
Memory benchmarks
|
Collaborator
Author
|
Pushing the commit that fixed the bugs spotted by @lostella caused all test jobs to fail because a new version of both Documenter and JET was released, which were stricter about some checks (or because the Julia version was changed to 1.13?) -- this is the reason why a couple of seemingly unrelated commits were also pushed. Anyway, @lostella, please let me know if the PR can be merged or if you need more time for further review. |
Member
|
@hakkelt this is now beyond my current context length limits (eheh), so I'll have Claude help out and report asap. |
lostella
requested changes
Oct 18, 2026
…t merge Adds AcceleratedDCTs test dependency, disables Aqua's persistent_tasks check, and carries forward the test coverage expansion, test refactors, and small dead-code/API cleanups developed on gpu-refactor after the GPU-support work was squash-merged upstream as 861af22.
It requires Julia >= 1.11 and is already added dynamically at runtime for that case in test/utils.jl; listing it statically forced dependency resolution to fail on the Julia 1.10 LTS CI job.
The private thread_count_functions registry and the @restrict_threading / @enable_full_threading macros in src/utils.jl are replaced by the new standalone NestedThreading.jl package. Beyond relocating the mechanism this fixes a latent correctness bug: the old save/restore was per scope, so two concurrent batched mul! calls captured each other's already-restricted thread counts and left BLAS permanently pinned at 1 with no scope active (reproducible before this commit, fixed after). NestedThreading refcounts budgets process-wide, snapshotting once and restoring once, and applies the minimum over all active scopes so nesting can only narrow a budget, never widen it. - src/utils.jl: delete the threading machinery. - SpreadingBatchOp/SimpleBatchOp: @restrict_threading @threads -> the @budgeted_threads macro; the LOCKING @sync/@Spawn sites use the function form with_restricted_threads. - FFTWOperators: drop the __init__ that reached into AbstractOperators' private registry; the FFTW pool now registers via NestedThreading's own package extension. - NFFTOperators: delete set_nfft_threading_expr and its two macros in favour of a with_nfft_threading helper over NestedThreading's scoped API, at all six call sites. - Inner libraries now get nthreads() / trip_count threads instead of always 1, so a 2-element batch on 8 threads runs BLAS at 4. - Add a regression test for the concurrency bug and for budget nesting. NestedThreading is wired in via [sources] pending registration in General, so these changes are not ready to push yet. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JHYpJyvm5XwhtMXLE1N8i9
BroadCast.jl's threaded adjoint mul! runs an outer @threads loop whose body calls mul! on arbitrary sub-operators, which may themselves use BLAS/FFTW. It never went through the old @restrict_threading, so it was a genuine oversubscription source that the migration had left in place. This was the last raw @threads in src/, so @threads is dropped from the imports; @Spawn and nthreads are still used. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JHYpJyvm5XwhtMXLE1N8i9
`TestUtils` is the setup module of nearly every testitem, and it unconditionally ran `GPUEnv.activate(; persist=true)` plus the FFTW-specific `AcceleratedDCTs` install. Any filtered run paid full GPU environment activation even when it selected no GPU test. Move both into a new `GpuEnvSetup` @testmodule (test/gpu_env_setup.jl). TestItemRunner evaluates setup modules lazily, only for testitems that survive the filter, so activation now happens iff a `:gpu` testitem is selected. All 47 `:gpu` testitems gain `GpuEnvSetup` in their setup list. `runtests.jl` publishes `ABSTRACTOPERATORS_TEST_GPU` from the filter closure, set from the tags of the items it actually accepts. Deriving it from the filter *string* would be wrong: GPU testitems also carry their category tags, so `:linearoperator` selects `Eye (GPU)` too. `GpuEnvSetup` errors when the flag is false, which can only mean a testitem missing its `:gpu` tag pulled the module in. Audit results: - `:gpu` testitems keep their inline `using GPUEnv`; the import is a name binding and costs nothing, only `activate` was expensive. - test_gpu_quality.jl's two items are tagged `:gpu` but exercise GpuExt through JLArrays alone, a direct test dep, so they deliberately do not use `GpuEnvSetup`. Measured: single non-GPU testitem 3m20s -> 8.7s. Pass counts unchanged -- `:jet` 224/224 and `:misc,:quality` 185/185 with zero GPU activation, `:calculus` 1557/1557 (still activates, correctly: it selects the `*(GPU)` items by their `:calculus` tag), `"Eye (GPU)"` 27/27 through the new setup module. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JHYpJyvm5XwhtMXLE1N8i9
Phases 1-4 of THREADING_PLAN.md. Threading was previously expressed three different ways with four scattered thresholds, and `copy_operator`'s deepcopy fallback silently ignored both of its keywords. Infrastructure (src/threading_policy.jl): - `is_threaded` trait, `supports_threading`, `default_threaded`/`threading_threshold` per-operator policy, and `adapt_operator`. - `copy_operator` now always copies; the "share when thread-safe" short-circuit moved to `adapt_operator`. The deepcopy fallback refuses `storage_type`/`threaded` instead of ignoring them, naming the type that needs a `_copy_operator_impl`. - Thresholds are transcriptions of benchmark/threading_sweep.jl (new), not guesses; each carries a PROVENANCE line. The sweep contradicted the plan's hypothesis that FastBroadcast wins or ties everywhere: for a pure copy `@.. thread=true` is time-identical to serial at every swept size, while `@batch` crosses over at 2^15. Nesting safety (the bug the plan called out): - create_BatchOp, create_threaded_SpreadingBatchOp and OperatorBroadCast adapted only the per-thread *copies*, leaving instance 1 threaded. All three now go through `_per_thread_operators`, which also branches on `is_thread_safe` to decide share vs copy -- `require_thread_safe` cannot manufacture thread safety, so it must not pretend to. - `opType` had to move after the adaptation: threading is a type parameter. Elementwise threading: all 10 nonlinear operators (SoftMax excluded, with the reduction reason recorded), FiniteDiff, and forwarding traits for the calculus operators. FiniteDiff `mul!` no longer allocates: `b[idx_1] .- b[idx_2]` materialised two temporaries per call. With `@views` it is allocation-free and 45x faster at n=1e6 (2581us -> 57us). Two problems the tests surfaced, both fixed in source rather than in the assertion: - `adapt_operator(op; threaded=false)` silently no-opped on forwarders, which lacked `is_threaded` and so inherited the `false` default while their children stayed threaded. - JET flagged runtime dispatch in `FiniteDiff(T, dims)`, which fell through to the `dir::Int` method and its `Val(dir)`. Split out a D=1 literal method. The two batching perf tests needed wider margins, not weaker ones: the FiniteDiff allocation fix sped the serial baseline up ~3x, shrinking an honest 2.5x speedup to 1.7x and letting runner noise flip a bare `<`. Both now take a minimum across repetitions at a workload size measured to favour batch parallelism (smaller-items/more-items was measured to hurt). Tests: :batching 283, :jet 245 (incl. new @test_opt/@test_call threading coverage), :quality 61, :calculus+:linearoperator+:nonlinearoperator+:misc 2698, :Threading+:Syntax 276 -- all passing. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JHYpJyvm5XwhtMXLE1N8i9
Phase 5 of THREADING_PLAN.md, for DCAT: its blocks are fully independent in both directions, with no accumulation and no shared output. The policy needed a second condition the plan did not anticipate. Aggregate work alone does not predict a win: 2 blocks x 2^18 elements is a 2^19 aggregate but measured only 1.03x, because two blocks cap the gain at 2x and the threading overhead eats it, while the same 2^19 aggregate across 8 blocks measured 1.9x. So `default_block_threaded` requires both `THRESHOLD_BLOCK_PARALLEL` (2^18, measured) and `MIN_BLOCKS_FOR_PARALLEL` (4). That pair admits every measured win and excludes every measured loss: nb=8 bs=2^16 1.90x threaded nb=8 bs=2^12 0.26x not threaded nb=4 bs=2^16 1.18x threaded nb=2 bs=2^18 1.03x not threaded nb=8 bs=2^18 9.67x threaded The threaded `mul!` is deliberately not `@generated`: a generated body must be pure and the threading macros expand to closures. The parallel loop is therefore a plain function and only the single-block body is generated, selected by `Val(i)`. That costs one dynamic dispatch per block -- nanoseconds against a whole child `mul!` -- and keeps the per-block type stability a closure over a heterogeneous tuple would have lost. `@budgeted_threads` rather than `@batch` because each body may itself reach BLAS/FFTW. Nesting safety applies here too: when the block loop threads, the constructor adapts the blocks to `threaded = false`. VCAT-forward and HCAT-adjoint from the plan are not included; they need their own measurement before shipping, per the plan's own rule that unmeasured variants are dropped. Tests: :Threading,:DCAT,:calculus 1769; :jet,:quality,:misc 462 -- all passing. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JHYpJyvm5XwhtMXLE1N8i9
…ression
Phase 6 of THREADING_PLAN.md, plus a regression this phase uncovered.
The regression: Phase 1's strict `copy_operator` fallback made batching *any* FFTW/DSP
operator raise. Those operators own scratch buffers, so a threaded batch needs one private
copy per thread; but they have no `_copy_operator_impl`, and the fallback refused the
`threaded = false` request the batch layer makes. The fix belongs in the fallback's rule,
not in per-operator patches: an operator with no threaded path cannot have a `threaded`
request violated, so deepcopy is a correct answer for it. That is the same rule
`_satisfies_constraints` already used, now applied consistently. `storage_type` is still
refused, since deepcopy genuinely cannot honour it.
Per subpackage:
- NFFTOp already carried a `threaded` field; it gains the traits and a `_copy_operator_impl`
that shares the immutable plan and dcf while giving the copy its own scratch buffer.
Requests that would need a different plan (thread count, storage backend) raise rather
than being silently ignored.
- DFT/IDFT gain `threaded` as the package-wide spelling of the existing `num_threads`, and
record the plan-time count so `is_threaded` can report it. FFTW is a counted pool, so this
is fixed at construction; switching it replans.
- DSP and Wavelet operators declare `supports_threading = false` explicitly rather than
inheriting the default, which is what lets a threaded batch wrap them.
Two bugs caught while writing this, both mine:
- A "cleanup" of a trailing `'` deleted the adjoint postfix from five IDFT constructors,
turning IDFT into a plain DFT and silently disabling the DFT/IDFT combination rules. The
FFTW combination tests caught it; there is now a test asserting IDFT stays adjoint-wrapped.
- The DFT copy replanned from the codomain element type, which is wrong for a real-input
DFT (`C == Complex{D}`) and produced an operator with a complex domain.
DCT/RDFT/IRDFT are not wired for plan-time thread counts -- they have no `num_threads`
plumbing at all today -- so they are declared unthreaded rather than half-wired. Recorded in
the plan's "Not done" section along with VCAT/HCAT block parallelism and the four legacy
thresholds that still gate the pre-existing operators.
THREADING_PLAN.md gains an execution record: measured-vs-guessed thresholds, the three plan
hypotheses that measurement contradicted, and what was deliberately left out.
Tests: full suite 3856 passed, 0 failed, 0 errored (21m47s).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JHYpJyvm5XwhtMXLE1N8i9
…all operators
Follow-up round covering VCAT/HCAT block threading, the remaining FFTW transforms, the
legacy thresholds, per-operator thresholds, and one semantics change that ties them
together.
`threaded` is now `Bool` only, with one meaning everywhere, resolved in one place
(`_resolve_threaded`):
false veto -- never threads. The only hard directive, and it must be: nesting
safety depends on a block/batch loop switching its children off and
being able to rely on it.
true (default) permission -- threads only if the policy also agrees.
Three families previously disagreed. Elementwise and block operators treated `true` as an
absolute override, batch operators as a permission, and the FFTW transforms bypassed the
policy entirely -- so a 256-point threaded DFT reported is_threaded == true while measuring
0.02x. `is_threaded` now reports what the operator will actually do, never what was asked
for; consequently a single-threaded session or GPU storage reports false without any
special-casing. `copy_operator`/`adapt_operator` keep accepting `nothing`, since there
`threaded` is a *constraint* ("no constraint, preserve") rather than a setting -- without it
a plain copy could not preserve an explicitly serial operator.
VCAT-forward and HCAT-adjoint now thread their block loops, and measuring them forced two
corrections to the Phase 5 policy:
- per-block, not aggregate: at a fixed 2^18 aggregate VCAT measures 1.32x with 4 blocks of
2^16 but 0.85x with 16 blocks of 2^14;
- per-operator: at 4 blocks of 2^16 VCAT measures 1.32x while HCAT measures 0.76x, because
HCAT's adjoint carries per-block ArrayPartition indexing. One shared constant would have
shipped that regression.
FFTW does thread r2r and r2c (measured 3.08x r2c, 2.25x DCT, 1.91x IDCT at n=2^22), so
DCT/IDCT/RDFT/IRDFT gained plan-time thread counts, and all six transforms are now gated on
measured crossovers (c2c 2^13, r2r 2^15, r2c 2^15). `num_threads` stays an explicit command
rather than a permission -- the escape hatch for callers who know their workload.
The four legacy thresholds are gone, which fixed two live defects:
- Variation's two constructors thresholded in different *units* (bytes vs elements), so the
same 10000-element input produced opposite threading;
- `_should_thread(::AbstractOperator)` had no size component at all, so batch operators
threaded four-element work.
Scale's legacy 1e4 cutoff was ~400x below its measured 2^22 crossover.
Thresholds are now per-operator transcriptions from benchmark/operator_thresholds.jl (new),
which measures each operator's real `mul!`. The cost classes proved too coarse: within
"transcendental" the crossovers span 2^8 (SoftPlus) to 2^11 (integer Pow), and Pow splits
internally by exponent kind since `x^0.5` lowers to `exp(p*log(x))`.
Also unified: SignAlternation had its own `threaded && nthreads() > 1` with no size policy
and no `is_threaded` method, so its trait always reported false regardless of what it did.
JET caught a regression that no functional test could: reassigning the captured `A` inside
the DCAT/VCAT/HCAT constructors boxed it to `Any`, cascading into 11 runtime-dispatch sites
in a constructor that should infer fully. Results were correct throughout; only @test_opt
saw it.
Tests: full suite 3971 passed, 0 failed. Three batching tests needed updating because they
encoded the old "threaded=true forces threading" assumption -- one of them was reading
`batch_op.operator[1]`, which on the single-threaded struct silently *slices* the wrapped
operator instead of erroring; it now goes through the `_wrapped_operator` accessor.
Known, unfixed, pre-existing: Variation's adjoint throws BoundsError when the trailing
dimension is exactly 2. Present on master; recorded in THREADING_PLAN.md.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JHYpJyvm5XwhtMXLE1N8i9
The adjoint's per-dimension branch chain tested `i == 2` before `i == size(y, d)`, so the last column of a dimension of length exactly 2 took the interior formula and read one element past the end. This hit *any* dimension of size 2, not only the trailing one -- (2,8) and (8,2,3) failed alongside (8,2) -- and the result was mathematically wrong there, not merely out of bounds: column n of the transpose is `y_1 + y_n` when n == 2, while the interior formula gives `y_1 + y_2 - y_3`. Rewrite the contribution as three independent terms (`_variation_adjoint_term`) rather than mutually exclusive branches, so the n == 2 case correctly takes both the `i == 2` and `i == n` terms. Both mul! methods now share one body so the threaded and serial paths cannot drift apart. `N` is passed to the shared body as a plain Int, deliberately not a Val: inside Polyester's `@batch` the loop body becomes a closure that does not carry the enclosing method's static parameters, so a `Val(N)` argument is built from a runtime value and every element dispatches dynamically. That allocated 293 MB at n = 2^22 and turned a 4.5x speedup into 0.25x, while leaving every correctness test green -- hence the new allocation testitem, timings being too noisy to assert on. Also reject singleton dimensions at construction. They have no finite difference to take and already failed with a BoundsError inside the *forward* kernel, so an ArgumentError naming the cause is strictly better; the array-based constructor now delegates to the tuple one so the guard and the threading policy each exist in one place. Verified against the dense forward matrix's exact transpose over 16 shapes x 2 element types x 2 threading modes. Threaded adjoint is back to baseline (12527 us vs 12668 us at n = 2^22, zero allocations) and the serial path is ~1.75x faster at mid sizes (3528 vs 6195 us at 2^18), the three-term form replacing unpredictable branches with straight-line arithmetic. Measured crossover is unchanged at 2^10. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JHYpJyvm5XwhtMXLE1N8i9
The plan is fully executed and its content is now recorded where it is actually used: threshold provenance in `threading_threshold` docstrings, policy rationale in `src/threading_policy.jl`, and the measurement methodology in `benchmark/operator_thresholds.jl`. Nothing references the file. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JHYpJyvm5XwhtMXLE1N8i9
The `Th` type parameter selected by the threading policy is consumed as `@.. thread = Th`, which requires a `FastBroadcast.True()`/`False()` instance, not the bare `True`/`False` type. Passing the type made every threaded elementwise `mul!` fail with `MethodError: no method matching DataType()`. Affects the nine elementwise non-linear operators plus `DiagOp` and `Scale`. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
For operators whose codomain and domain shapes have different types the `size(L)` tuple is heterogeneous, so indexing it with a runtime `i` widens the result to a `Union`. That `Union` propagated through `ndims`, `count_dims`, `ndoms` and `check`, turning every comparison in the `mul!` argument check into a runtime dispatch under `@test_opt`. Index the size tuple with literal indices at each of those sites instead: `ndims` maps over `size(L)`, `ndoms` recomputes from `size(L)` rather than mapping over `ndims(L)`, and `check` destructures the tuple once. `size(L, i)` itself is public API and keeps its runtime-index method. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`y[range, d]` builds a `SubArray` whose index tuple mixes a `StepRange` and an `Int`. Base's `ensure_indexable` maps over that tuple, and abstract interpretation widens the per-element type, which left the resulting view's index type non-concrete and made every broadcast in the kernel a reported runtime dispatch. Compute the column offset explicitly and index `y` with a single range instead, in both the threaded and the non-threaded method. Behaviour is unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This reverts commit e7753ac. The "failed to optimize due to recursion" reports it worked around were downstream of the non-concrete `size`/`ndims`/`ndoms` chain, now fixed at the root. With that fix in place the `@generated size` methods infer fine while binding only the type parameters they use, and the full `:jet` slice passes. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
An explanatory comment between the docstring and `ndims(L::AbstractOperator)` detached the docstring from the method, so Documenter failed the build with `no docs found for 'ndims' in @docs block in docs/src/properties.md`. Move the comment above the docstring. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`_fbbool(::Type{FastBroadcast.True})` accepted the bare type, so an operator
that stored the type instead of the singleton instance still reported the right
`is_threaded` trait and only failed later, inside `@.. thread = Th`, with
`MethodError: no method matching DataType()`. No caller passes a type any more;
removing the methods turns that mistake into an immediate `MethodError` at the
trait instead of a silent mis-encoding.
Replace the two tests that pinned the type form with a check that the operators
which carry the flag really store an instance.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
AGENTS.md allows only a fixed set of lowercase type tags and CamelCase operator names, so :Threading was neither. Rename it to :threading and add that tag to the documented list. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Four review findings on the block-parallel loops: - Data race: a threaded block loop called `mul!` on the same operator from several tasks when one non-thread-safe operator filled several blocks, directly (`DCAT(C, C)`) or wrapped (`DCAT(C, 2 * C)`), so the blocks shared its scratch buffers and the output was wrong on every call. Each non-thread-safe block now gets its own copy (`_block_loop_operator`). - Storage: only the first block's storage was checked, so a CPU block in front of device blocks threaded the loop. Every block is checked now; a multi-domain block counts as CPU when all its components are. - Lost flag: `copy_operator(dcat; threaded = true)` forced the blocks serial even when the constructor then left the block loop serial, and `get_normal_op(::DCAT)`, `VCAT.remove_slicing`, `getindex`, `Jacobian` and the combination rules rebuilt block operators without their block-loop flag, so a `threaded = false` veto was lost. `HCAT(A, buf)` and `VCAT(A, buf)` take `threaded` for that. - `check`: HCAT's `Tuple` `mul!` methods (forward, and serial and threaded adjoint) did not call `check`, so a wrong size surfaced as a block's error, or a `TaskFailedException` from the threaded loop. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- `copy_operator(DiagOp(T, dims, 2.0))` threw a `MethodError` from `_array_wrapper(::Float64)`. A scalar diagonal now keeps the operator's own storage type. - Copying an `LBFGS` built a fresh operator and dropped the curvature history (`s_M`, `y_M`, `ys_M`, `currmem`, `curridx`, `H`), so the copy acted as the identity. The history is the operator's state: the copy now starts from it, with its own buffers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- FFTW and DSP operators on device storage reported `is_threaded == true` above the size gate: `_fftw_num_threads` and `_dsp_fftw_num_threads` never looked at storage, so `RDFT`, `IRDFT`, `Conv` and `Xcorr` on a `CuArray` recorded `Threads.nthreads()` and `adapt_operator(op; threaded = false)` replanned them on every call. Both now return 1 for device storage, whatever was asked. - Device-backed `MatrixOp`/`LMatrixOp` entered the restricted BLAS scope (a process-wide lock plus a setter round trip for every pool) on every `mul!`, although a device `gemm` uses none of those pools. The scope is now skipped for non-CPU storage: 2x2 `JLArray` `MatrixOp` `mul!` 5.44 us and 64 B before, 0.073 us and 0 B after (bare `mul!`: 0.064 us); the CPU `threaded = false` path is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`==` on SimpleBatchOp and SpreadingBatchOp compared the wrapped operators and the sizes but not the domain/codomain batch masks, so two batch operators batching along different axes compared equal while giving different results. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`threaded = true` only permits threading, and with `-t 1` it correctly resolves to serial, but the threading contract helper and the threading traits tests asserted `is_threaded(...) == true` unconditionally. CI only passed because it sets `JULIA_NUM_THREADS: 2`. Those assertions now expect `Threads.nthreads() > 1`; none is removed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
It was the one :gpu testitem still on `setup = [TestUtils]` only. Without `GpuEnvSetup`, `gpu_backends()` returns nothing and the testitem passes without asserting anything, or not, depending on test order. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
copy_operator, get_normal_op, permute, remove_displacement and Jacobian passed blocks that share no scratch through the constructor, which copied the non-thread-safe ones again. copy_operator(D; threaded = true) also built threaded block copies (FFTW plans) only to replace them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The JLArray gemm allocates by itself, so the test failed on CI. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Batch operators fell back to `supports_threading == false`, so `adapt_operator(b; threaded = false)` returned a threaded batch op unchanged. Their copies also dropped `threaded` for the wrapped operator, while a serial batch loop reported `is_threaded == false` even around a threaded operator. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Copying every block that is not thread-safe detached the loop's blocks from the caller's objects, so e.g. `update!` on an `LBFGS` block no longer reached a threaded DCAT. Copy a block only when it shares a non-thread-safe operator with another block. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`is_threaded` and `supports_threading` now default to asking the operator's children, so the forwarders no longer repeat the same two methods and a new wrapper cannot forget them. Also: - `_block_loop_operators` walks each block's tree once. - `_children(::SpreadingBatchOpCopying)` iterates instead of allocating. - `_should_thread(op, storage_type)` takes `storage_type` as a default argument. - `DiagOp` copy uses `_fbbool`. - `DFT`/`IRDFT` copies with nothing to change return the operator itself. - Batch operator docstrings no longer mention a `nothing` default for `threaded`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
hakkelt
force-pushed
the
unified-threading
branch
from
November 3, 2026 14:16
117b690 to
88fd54f
Compare
- `is_threaded` and `supports_threading` referred to the internal `_children` with an `@ref` that Documenter cannot resolve, which fails the docs build. - The `_children` docstring still promised a tuple and described the old `false` default. - The `SimpleBatchOp` comment described a removed method. - The `DFT`/`IRDFT` copy comments now say why returning `op` is a copy. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Benchmark Results (Julia v1.13.3)Run once per thread count ( 1 thread🚀 6 benchmarks improved in time · 🚀 7 benchmarks use less memory Time benchmarks
Memory benchmarks
2 threads🚀 25 benchmarks improved in time · 🐢 1 time regression detected · 🚀 6 benchmarks use less memory Time benchmarks
Memory benchmarks
|
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The `update!` call used random gradients, so the curvature check could reject the pair and leave the operator an identity map that any stale copy would also match. Use `g = 2x`, assert that the update changed the operator, and compare with `≈`: inside the block loop BLAS runs single-threaded and rounds the LBFGS dot products differently, which failed the exact comparison on Julia 1.10. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…wdown - A literal power in a threaded `@..` broadcast passes `^` and the exponent in `Ref`s, which cost 2 allocations per call in Sigmoid forward and the Tanh and Atan Jacobian-adjoints. Use `inv` or a scalar helper instead. - The fused Sigmoid Jacobian-adjoint was ~1.3x slower than master on one thread. Computing `exp` in its own pass first is faster both serial and threaded (32768 elements: 264 -> 187 us serial, 146 -> 102 us on 2 threads). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Collaborator
Author
Well, that's fully understandable given the size of this PR. 😄 I reviewed it manually before I submitted, but it seems that either I wasn't careful enough or it is beyond the capabilities of a human reviewer to catch all bugs at this scale... Anyway, I generated fixes for the spotted bugs and reviewed the quality of the changes and whether they fix the problematic behavior both manually and by AI. Some other bugs were found and fixed along the way, and I merged two minor PRs that upgraded two dependencies (GPUEnv and JET), so that's why this PR was rebased on master. Btw, do you agree with the proposed API, @lostella? The |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR unifies the interface for enabling/disabling threading across operators and adds multithreading support to all possible operators, replacing scattered per-operator implementations with a consistent policy infrastructure.
Important Changes
Threading Policy Infrastructure
src/threading_policy.jlmodule introducing a unified threading framework:is_threadedtrait for operators with optional threading supportsupports_threading,default_threaded, andthreading_thresholdper-operator policiesadapt_operatorandcopy_operatorfunctions for consistent threading controlbenchmark/threading_sweep.jl) replace previous guessesOperator Threading Coverage
Elementwise nonlinear operators (9 operators: Pow, Exp, Sin, Cos, Atan, Tanh, Sech, Sigmoid, SoftPlus)
FiniteDiff optimization: eliminated allocations from
b[idx_1] .- b[idx_2]materialization using@viewsCalculus operators (Sum, HCAT, VCAT, DCAT, Compose, Scale, BroadCast, etc.) add forwarding traits
Batching operators (BatchOp, SimpleBatchOp, SpreadingBatchOp)
Bug Fixes
_per_thread_operators(branch onis_thread_safe)Test Expansion
@test_optand@test_callfor threading APIsInfrastructure & Cleanup
Performance Baselines
benchmark/threading_sweep.jl) captures 243 measurements for future regression detection