GPU Refactor: Add array_type parameter and GPU extensions#43
Draft
hakkelt wants to merge 50 commits into
Draft
Conversation
Add array_type keyword constructors and domain_storage_type/codomain_storage_type storage type traits to all operators. Add core GPU extension (GpuExt) with operator overrides for GetIndex, Variation, and ZeroPad. Add GPU extensions for DSPOperators, FFTWOperators, NFFTOperators, and WaveletOperators subpackages. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Update test infrastructure to support GPU (JLArray) testing. Add :jlarray tags to relevant testitems. Add gpu_utils.jl helper. Add GPU quality tests. Update operator testitems with proper tags (:linearoperator, :nonlinearoperator, etc.). Rename CpuWrapper tests to OperatorWrapper. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Updated tests for Eye, FiniteDiff, GetIndex, L-BFGS, LMatrixOp, MatrixOp, MyLinOp, Variation, Zeros, and nonlinear operators to utilize GPUEnv for backend management. - Removed specific CUDA and AMDGPU checks, replacing them with a loop over available GPU backends. - Simplified test setups by eliminating redundant code and ensuring compatibility with various GPU array types. - Ensured all tests are now tagged appropriately for GPU execution without dependency on specific GPU libraries.
- Updated various test files to replace domain_storage_type and codomain_storage_type with domain_array_type and codomain_array_type for consistency and clarity. - Removed unnecessary verbose print statements in tests to streamline output. - Adjusted GPU-related tests to ensure proper handling of array types. - Ensured that all tests maintain functionality while improving readability and maintainability.
…y backend-specific limitations
…n storage type parameters
… documentation CI job
Co-authored-by: Copilot <copilot@github.com>
Replace the FFT-based adjoint mul! with a tiled FIR direct convolution on
CPU paths (H <: Array{T}). The GPU fallback keeps the FFT-based approach.
Algorithm: y[j] = Σ_k h[k] * b[padlen+j-k], unrolled 8-wide so all
accumulators (a0..a7) live in registers. Reads b[base:base+7] consecutively
per inner k-iteration (cache-friendly), writes y only once.
Benchmark (n=32768, h length 21, Float64, 1 FFTW thread):
- Before (65536-pt FFT): ~484 μs
- After (tiled FIR): ~107 μs (~4.5× speedup)
Baseline on benchmark machine was ~418 μs, so this should close the
regression seen in PR kul-optec#40.
- Add XcorrAdjFFT helper struct carrying the adjoint FFT buffers and plans
- Xcorr.adj_fft is Nothing for CPU arrays (H <: Array) — no adjoint FFT
buffers are allocated; the tiled FIR path (mul! on Xcorr{<:Array}) is used
- For GPU backends adj_fft is a XcorrAdjFFT; the FFT-based adjoint dispatches
on <:XcorrAdjFFT instead of the old top-level struct fields
- Add _xcorr_plan_kwargs helper: passes flags=FFTW.MEASURE for CPU Arrays only,
no flags for GPU backends — fixes FFTW.MEASURE incompatibility with cuFFT
Add array_type keyword constructors and domain_storage_type/codomain_storage_type storage type traits to all operators. Add core GPU extension (GpuExt) with operator overrides for GetIndex, Variation, and ZeroPad. Add GPU extensions for DSPOperators, FFTWOperators, NFFTOperators, and WaveletOperators subpackages. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
… docs Add check() calls to all mul! methods for consistent argument validation. Fix dispatch ambiguities in HCAT/VCAT/DCAT/BatchOp. Fix IRDFT storage type. Fix SignAlternation GPU broadcast. Fix return-if pattern in utils.jl. Rename CpuOperatorWrapper to OperatorWrapper in calculus. Add GPU crossover benchmarks. Update documentation for GPU support and operator properties. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Update test infrastructure to support GPU (JLArray) testing. Add :jlarray tags to relevant testitems. Add gpu_utils.jl helper. Add GPU quality tests. Update operator testitems with proper tags (:linearoperator, :nonlinearoperator, etc.). Rename CpuWrapper tests to OperatorWrapper. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
FiniteDiff (GPU), DiagOp (GPU), Eye (GPU), and MyLinOp (JLArray) all use JLArrays internally but were tagged :gpu without :jlarray. This means they ran under the general :gpu filter but were excluded from the :jlarray-only GPU coverage runs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Remove dead @generated branches in VCAT.jl: the Vararg constructor always flattens inner VCATs so P type params are always Int, not tuples. Removed _ndoms_from_type(VCAT, dim), simplified @generated constructor, @generated mul! (forward and adjoint), @generated size and @generated codomain_type. - Add tests for previously uncovered Compose paths: - Type mismatch error (DomainError) in constructor - L'*L shortcut returning get_normal_op - get_normal_op(Compose) if branch (has_optimized_normalop) - Scale(1, Compose) identity early-return - Adjacent adjoint optimized normal GetIndex*GetIndex' - _copy_operator_impl via copy_operator - remove_slicing is_eye path (AffineAdd wrapping GetIndex, length 2 and 3) - Add tests for previously uncovered combination_rules.jl paths: - Scale+Compose forwarding (line 82), AdjointScale+Compose (line 90) - Compose+Scale forwarding (lines 98-99), Compose+AdjointScale (lines 109-110) - Scale+MatrixOp, Scale+AdjointMatrixOp, AdjointScale+DiagOp - DiagOp+Scale, AdjointDiagOp+Scale forwarding - Add DCAT nested DCAT test: exercises @generated tuple-index branches (lines 81-83, 120-121, 167-168, 200-202, 271-273) - Add VCAT constructor error tests (DimensionMismatch, type mismatch) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…e{N,Int}
The Vararg constructor unconditionally flattens inner VCATs, so idxs
is always NTuple{N, Int}. The old comment describing the (1,(2,3)) case
was hypothetical and unreachable via the public API.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Add N isa Int guard to @generated VCAT(NTuple, buf) to prevent MethodError when JET analyzes with N=TypeofVararg - Replace remove_slicing.(generator) with collect(map(..., L.A)) to avoid the L[i] VCAT-indexing path that creates Tuple{Vararg} and triggers the @generated constructor with unknown N Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Fix _is_dft_op to correctly distinguish domain vs codomain shifts
when the operator is a Compose. Previously, the function used
'size == :domain' (always false) and a permissive all() check that
allowed diagonal operators on either end regardless of shift side.
Correct semantics:
- Domain shift: innermost (first) op must be DFT-like; all outer
ops must be diagonal (they commute with SignAlternation).
- Codomain shift: outermost (last) op must be DFT-like; all inner
ops must be diagonal (they commute with SignAlternation).
This prevents incorrectly applying the SignAlternation optimization
to cases like D*DFT with codomain shifts or DFT*D with domain shifts.
- Add tests for _check_shift_dirs error branches:
- Too many dirs (N < M)
- Unsorted dirs
- Add tests for _is_dft_op with Compose operators:
- D*DFT with domain shift (valid optimization → SignAlternation used)
- D*DFT with codomain shift (optimization not valid → FFTShift used)
- MatrixOp*DFT with domain/codomain shifts (no optimization)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The @generated function VCAT(AA::NTuple{N,...}, buf) was crashing with SIGILL ("Unreachable reached") when called with a Tuple{Vararg{...}} (N not statically known). This is a Julia 1.12 bug: returning a @NoReturn expression (error()) from a @generated specialization for TypeofVararg causes the JIT to emit an unreachable instruction that is then hit at runtime. Fix: add an else branch that returns a runtime-length fallback instead of an error call -- ntuple(identity, length(AA)) builds the idxs tuple at runtime, avoiding the SIGILL. Also change Zeros.jl multi-codomain constructor to use ntuple(f, Val(NN)) instead of a vector comprehension + splat, so the tuple length is always statically known and the fast @generated path is used. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…_slicing second branch
- Remove dead overload _to_gpu_indices(AbstractGPUArray, AbstractGPUArray{<:Integer})
in ext/GpuExt/linearoperators/getindex.jl (no call sites)
- Add GPU getindex tests: all-indices→Eye shortcut, GPU boolean mask,
full-mask→Eye shortcut (covers lines 19-20, 35-44, 46-55)
- Add ext/GpuExt/properties.jl coverage: _should_thread(instance) and
storage_type_display_string assertions in test_gpu_quality.jl
- Add test for VCAT.remove_slicing second branch (elseif path):
VCAT of HCATs with different-sized GetIndex operators triggers the
domain-size fixup logic (lines 174-198 in VCAT.jl)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ed properties - Test Variation(jl(...); threaded=true) on JLArrays to cover GPU threaded mul! dispatch paths (variation.jl lines 92-94, 104-108) - Add SimpleBatchOpMultiThreaded properties testitem that directly constructs the multi-threaded variant (bypassing nthreads() > 1 guard) to cover diag_AAc, diag_AcA, diag, has_optimized_normalop, opnorm, estimate_opnorm methods (SimpleBatchOp.jl lines 285-368) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Cover the two remaining uncovered branches of the generic combine(L, R) fallback (properties.jl lines 345, 354): - combine(L, Eye) where L is not Eye → returns L - combine(L, Zeros) where L is linear, square, zero displacement → returns Zeros Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Cover previously uncovered lines in: - WaveletOperators: constructor errors for odd dimensions and excessive levels - DSPOperators: Xcorr complex domain FFT plan path - FFTWOperators: DFT ORTHO normalization, Shift out-of-range dirs error - src/properties.jl: AbstractOperator fallback property methods - src/calculus/: AffineAdd type mismatch, Ax_mul_Bxt/Axt_mul_Bx 2D DimensionMismatch, BroadCast same-size return and non-compact adjoint reshape, Compose combination paths (lines 73/81-82/88), HCAT fun_name reversed and get_slicing_expr Tuple-of-Tuple path, Sum all-Zeros case, VCAT remove_slicing error, Scale+Compose combination else branches - src/linearoperators/MatrixOp.jl: Scale(coeff, AdjointMatrixOp) paths - src/batching/SpreadingBatchOp.jl: BatchOp without sizes, unsupported threading strategy error - ext/LinearMapsExt.jl: isposdef delegation Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Remove dead VCAT getindex error branch (VCAT.idxs always NTuple{N,Int})
- Remove dead Compose getindex diagonal-tail branch (combination rules always
simplify diagonal*HCAT to HCAT, making this branch unreachable)
- Remove dead properties.jl:345 branch (is_eye(R) in generic combine shadowed
by combine(T1, ::Eye) = T1 in combination_rules.jl)
- Fix HCAT getindex dead code test → replace with tuple-idxs error path test
- Fix test_combinations.jl: correct FiniteDiff*MatrixOp dimension order
- Add tests for: Sum multi-domain getindex (syntax.jl:71), Scale multi-domain
getindex (syntax.jl:132), check() ArrayPartition errors (utils.jl:104/128),
Scale complex coeff on real AdjointMatrixOp (MatrixOp.jl:84), HCAT tuple-idxs
error path (syntax.jl:87), Compose multi-domain error branch (syntax.jl:62)
- syntax.jl now at 100% coverage; overall coverage 99.0%
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01W3roEtQf9X1mmFj3epwNcR
The be7dce0 commit partially renamed the storage-type trait functions (domain_array_type → domain_storage_type) throughout the source and tests, but left the public exports, internal callers in properties.jl, and NffftNormalOp.jl still using the old names, making the package unloadable (MethodError on domain_array_type for NFFTOp and similar). Also left duplicate _array_type_for_elem / _array_wrapper_type / _array_wrapper / _normalize_array_type definitions causing "Method overwriting is not permitted during Module precompilation". This commit completes the rename uniformly: - `domain_storage_type` → `domain_array_type` across all src/, test/, and NFFTOperators/src/ files - Removes the duplicate helper functions from properties.jl - Removes the erroneous self-alias lines introduced by a previous fix attempt Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YAfb7wBXYQDTwoqPvwViP9
…rename - DCAT: replace allocating vcat/invperm in _dcat_apply_invperm with tuple-recursive helpers (_dcat_flatten_idxs, _dcat_find_in_tuple) eliminating 60-260 heap allocations per mul! call; also use _storage_eltype consistently in _compute_dcat_ds/cs - Xcorr: restore CPU tiled 8-wide FIR adjoint dispatch (30× faster than FFT path on CPU); keep GPU FFT adjoint unchanged; restore _xcorr_plan_kwargs (FFTW.MEASURE) for forward and adj plans - GpuExt: rename storage_type_display_string → array_type_display_string to match the properties.jl rename; fixes precompilation error and GpuExt JET test failure - Test: fix "MIMOFilt: GPU constructor" test item name (was duplicate) and add GpuTestUtils to setup so jl() is available; add :jlarray tag Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YAfb7wBXYQDTwoqPvwViP9
… items - Filt: GPU constructor test was missing GpuTestUtils (jl undefined) and DSP import (filt undefined); renamed to avoid duplicate name with the non-GPU properties test; added :jlarray tag - MIMOFilt: GPU constructor test was missing DSP import for filt(); the jl undefined error was masking this in the previous run Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YAfb7wBXYQDTwoqPvwViP9
Applies changes from dangling commit 15cce9d (dropped during rebase). Replaces per-device JLArrays/CUDA/AMDGPU-specific test items with generic `for backend in gpu_backends()` loops using GPUEnv helpers (gpu_randn, gpu_zeros, gpu_wrapper, to_gpu). Also removes GpuTestUtils testmodule and adds GPUEnv.activate() to the shared TestUtils setup. CPU-only test items added on the gpu-refactor branch are preserved.
…8692)
- OperatorWrapper: remove ndims from storage type parameters (S{T} not S{T,N})
- test_operatorwrapper: update assertions to Array{Float64} (without ndims)
- Quality tests: add persistent_tasks=VERSION>=v"1.11"; drop now-unnecessary
stale_deps=false and ambiguities=false flags from fftw/nfft/wavelet Aqua tests
- test/utils.jl: move AcceleratedDCTs install after GPUEnv.activate
- test/Project.toml: remove AcceleratedDCTs from explicit deps
- docs/src/gpu.md: rename storage_type→array_type, add backend-specific notes
section, add DCT/IDCT AcceleratedDCTs section, improve CpuOperatorWrapper docs
- docs/src/properties.md: remove stale has_mutable_buffers entry
- .github agents/instructions/skills: rename storage_type→array_type
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01YAfb7wBXYQDTwoqPvwViP9
…HCAT/VCAT/Sum Several GPU test items called GpuBackend instances directly as conversion functions (e.g. `backend(x)`), but GpuBackend has no callable method; replace with `to_gpu(backend, x)`. Also fix world-age-unsafe manual allocation in the nonlinear operator GPU test helper by allocating via the operator's own `*`/adjoint `*` instead of `similar(x, T, dims...)`, and construct several nonlinear/linear GPU test operators via their type-based constructors with an explicit `array_type` rather than from a GPU array instance, avoiding stale-buffer/storage-type mismatches. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The check is unreliable and not relevant for this package. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
hakkelt
marked this pull request as draft
July 17, 2026 16:35
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.
This PR introduces a comprehensive GPU refactoring.
array_typeparameter to support multiple backends.GPUEnvfor backend-agnostic testing.