Skip to content

GPU Refactor: Add array_type parameter and GPU extensions#43

Draft
hakkelt wants to merge 50 commits into
kul-optec:masterfrom
hakkelt:gpu-refactor
Draft

GPU Refactor: Add array_type parameter and GPU extensions#43
hakkelt wants to merge 50 commits into
kul-optec:masterfrom
hakkelt:gpu-refactor

Conversation

@hakkelt

@hakkelt hakkelt commented Jul 17, 2026

Copy link
Copy Markdown
Collaborator

This PR introduces a comprehensive GPU refactoring.

  • Adds the array_type parameter to support multiple backends.
  • Introduces new GPU extensions for various operators.
  • Modernizes the test infrastructure using GPUEnv for backend-agnostic testing.
  • Improves test coverage across nonlinear operations and handles several edge cases.

hakkelt and others added 30 commits April 27, 2026 16:01
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.
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>
hakkelt and others added 20 commits June 29, 2026 19:04
…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
hakkelt marked this pull request as draft July 17, 2026 16:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant