GPU override coverage: tracking and roadmap

[zazabap]

Following Manifolds.jl#856, CUDA overrides live in this package. This issue tracks progress across manifolds to coordinate contributions.

All overrides target PowerManifold{ℝ, <:M, <:Tuple, ArrayPowerRepresentation} with CuArray{T,3}, replacing ManifoldsBase's sequential get_iterator loop with batched GPU ops. For many operations the single-matrix CPU code already works on CuMatrix (e.g. svd, /), but the PowerManifold loop launches N separate kernels — overrides batch these into single strided calls.

Shared helpers (helpers.jl)

• _matrix_exp_gpu — batched matrix exponential (Taylor + scaling-and-squaring)
• _matrix_log_gpu — batched matrix logarithm (inverse scaling-and-squaring + Denman-Beavers)
• _batched_inv_gpu — batched matrix inverse (LU)
• _batched_sqrtm_gpu — batched matrix square root (Denman-Beavers)
• _cholesky_qr_gpu! — batched Cholesky-QR (A'A + potrfBatched! + trsm_batched!) (PR add GPU retract_qr_fused! via Cholesky-QR #11)

GeneralUnitaryMatrices (PR #4)

Default retraction: ExponentialRetraction() — already overridden. Default vector transport: ProjectionTransport() — already overridden. Both default paths work on GPU.

• exp! — _matrix_exp_gpu + gemm_strided_batched (generic-n; covers n=2,3,4 specializations)
• log! (ℝ and ℂ) — _matrix_log_gpu + skew-symmetrization
• inner, norm
• project! (point, AbsoluteDeterminantOneMatrixType) — gesvdj!. Note: Rotations (DeterminantOneMatrixType) needs additional det correction — not covered.
• project! (tangent) — skew(p'X) via batched GEMM
• retract_polar_fused! — gesvdj! with CPU fallback for large matrices
• retract_qr_fused! — Cholesky-QR via _cholesky_qr_gpu!; fully batched, no size limit (PR add GPU retract_qr_fused! via Cholesky-QR #11)

Nice-to-have

• parallel_transport_to! — building blocks exist but not wired for batched PowerManifold
• rand! — scalar indexing in CPU code; can initialize on CPU and transfer
• inverse_retract_polar! — Rotations version uses lyap() + try/catch; no GPU path

Stiefel (PRs #12, #11)

Default retraction: PolarRetraction() — already overridden. Default vector transport: DifferentiatedRetractionVectorTransport(PolarRetraction()) — does not use log! or parallel_transport_to!. Only ℝ field covered; complex Stiefel has no overrides yet.

• exp! (ℝ) — block-diagonal 2k×2k _matrix_exp_gpu + batched GEMM
• inner, norm
• retract_polar_fused! — gesvdj! with CPU fallback for large matrices
• project! (point) — _polar_project_gpu! via batched SVD (PR add GPU project! (point + tangent) for Stiefel #12)
• project! (tangent) — X - p*sym(p'X) via batched GEMM (PR add GPU project! (point + tangent) for Stiefel #12)
• retract_qr_fused! — Cholesky-QR via _cholesky_qr_gpu! (PR add GPU retract_qr_fused! via Cholesky-QR #11)

Nice-to-have

• log! — delegates to iterative shooting via StiefelSubmersionMetric; hard to GPU-ify
• parallel_transport_to!
• inverse_retract_polar! — needs lyap(); no GPU equivalent
• inverse_retract_qr! — scalar loop with @inbounds
• rand! — can initialize on CPU and transfer

Grassmann (PRs #6, #7, #8, #11)

Default retraction: ExponentialRetraction() — overridden in PR #6. Default vector transport: ParallelTransport() — overridden in PR #8. CG/L-BFGS optimizers on Grassmann now have full GPU support. Only ℝ field covered.

• inner, norm
• project! (point) — _polar_project_gpu! via batched SVD (PR add GPU project! and retract_polar_fused! for Grassmann #7)
• project! (tangent) — X - p*(p'X) via batched GEMM (PR add GPU project! and retract_polar_fused! for Grassmann #7)
• retract_polar_fused! — delegates to _polar_project_gpu! (PR add GPU project! and retract_polar_fused! for Grassmann #7)
• exp! — SVD + polar orthogonalization, replaces CPU's qr(z).Q round-trip (PR add GPU exp! for Grassmann, gesvda! fallback, benchmark cleanup #6)
• log! — batched inverse (_batched_inv_gpu) + gesvdj! + atan.(S) (PR add GPU log!, inverse_retract_polar!, parallel_transport_direction!, distance for Grassmann #8)
• inverse_retract_polar! — q * inv(p'q) - p via batched LU inverse + in-place GEMM (PR add GPU log!, inverse_retract_polar!, parallel_transport_direction!, distance for Grassmann #8)
• parallel_transport_direction! — SVD of direction + sin/cos rotation + projection (PR add GPU log!, inverse_retract_polar!, parallel_transport_direction!, distance for Grassmann #8)
• distance — inverse retract → batched SVD → norm(atan.(S)) (PR add GPU log!, inverse_retract_polar!, parallel_transport_direction!, distance for Grassmann #8)
• retract_qr_fused! — Cholesky-QR via _cholesky_qr_gpu! (PR add GPU retract_qr_fused! via Cholesky-QR #11)

Nice-to-have

• rand! — can initialize on CPU and transfer

Sphere

• inner, norm
• Most other ops are already GPU-safe in Manifolds.jl (broadcasting + dot). PowerManifold of Sphere ≈ Oblique — low priority.

Known limitations

• gesvdj! size limit: cuSOLVER's batched Jacobi SVD fails beyond a supported matrix size. Existing overrides have try/catch CPU fallback.
• Real field only: Stiefel and Grassmann overrides dispatch on T <: Real. Complex manifolds have no batched overrides yet (GeneralUnitaryMatrices handles both via T <: Number).

Not planned (fundamentally hard)

• SPD: nearly all ops need eigen() multiple times
• Rotations small-n (n=4): scalar indexing + try/catch (generic-n override covers this)
• Hyperbolic: core ops are GPU-safe; minkowski_metric scalar indexing in some paths

Proposed contribution order

1. Grassmann exp! — default retraction, effectively broken for GPU users ✓ PR #6
2. Grassmann project! (point + tangent) + retract_polar_fused! — enables polar retraction as alternative ✓ PR #7
3. Grassmann log! + inverse_retract_polar! — needed by distance and parallel_transport_to! ✓ PR #8
4. Grassmann parallel_transport_direction! + distance — needed for CG/L-BFGS and convergence checks ✓ PR #8
5. Stiefel project! (point + tangent) ✓ PR #12
6. retract_qr_fused! for GeneralUnitaryMatrices — Cholesky-QR approach ✓ PR #11 (also covers Stiefel + Grassmann)
7. Specialized overrides for Rotations n=2,3 — generic-n works but specialized CPU variants are much faster (per @mateuszbaran's feedback in this thread)

Each as a separate PR with JLArray + CUDA tests. I'd love to hear your suggestions on priorities, approaches, or anything I may have missed or gotten wrong — any feedback from the experts here would be greatly appreciated!

Related

• Manifolds.jl#856 — original PR and discussion
• #4 — merged GeneralUnitaryMatrices PR
• #6 — merged Grassmann exp! PR
• #7 — merged Grassmann project! + retract_polar_fused! PR
• #8 — merged Grassmann log!, inverse_retract_polar!, parallel_transport_direction!, distance PR
• #11 — Cholesky-QR retract_qr_fused! for GeneralUnitaryMatrices, Stiefel, Grassmann PR
• #12 — merged Stiefel project! (point + tangent) PR
• #2 — closed Euclidean PR
• 15da7e3 — Stiefel exp! sketch by @mateuszbaran

[zazabap]

For the Cholesky-QR approach mentioned for retract_qr_fused! and Grassmann exp!: the key reference is CholeskyQR2 (Fukaya & Nakatsukasa, 2014) — run CholQR twice for stability. This is a different algorithm from Manifolds.jl's Householder QR (qr() → orgqr!), chosen because all primitives are batchable on GPU (gemm_strided_batched for A'A, potrfBatched! for Cholesky). For retraction where A ≈ I + small step, even single-pass CholQR should be well-conditioned.

[mateuszbaran]

Thanks for making a roadmap, I think it's good to have it. I have recently added docs here with a benchmark: https://juliamanifolds.github.io/ManifoldsGPU.jl/dev/ and as work progresses it would be nice to expand it as the work progresses. The table can be regenerated using the benchmark script here: https://github.com/JuliaManifolds/ManifoldsGPU.jl/blob/main/benchmarks/main.jl .

Rotations small-n (n=2,3,4): scalar indexing + try/catch (generic-n override covers these)

PowerManifold of Rotations for n=2 and 3 could be fairly important in practice, and scalar indexing shouldn't be an issue there. Generic n code is much slower than these specialized variants.

SPD is indeed a bit problematic, so leaving it as "maybe later" makes sense.

Proposed contribution order

1. Grassmann exp! — default retraction, effectively broken for GPU users
2. Grassmann project! (point + tangent) + retract_polar_fused! — enables polar retraction as alternative
3. Grassmann log! + inverse_retract_polar! — needed by distance and parallel_transport_to!
4. Grassmann parallel_transport_direction! + distance — needed for CG/L-BFGS and convergence checks
5. Stiefel project! (point + tangent)
6. retract_qr_fused! for GeneralUnitaryMatrices — Cholesky-QR approach

Each as a separate PR with JLArray + CUDA tests. I'd love to hear your suggestions on priorities, approaches, or anything I may have missed or gotten wrong — any feedback from the experts here would be greatly appreciated!

I think this is a good plan. It would also be nice to add n=2 and n=3 Rotations at some point.

Regarding CG/L-BFGS, they need some kind of vector transport (not necessarily parallel) and AFAIK there isn't a clear consensus on which transport works the best. Projection VT often works well enough, and some argue that differentiated retraction is better suited than parallel transport. Anyway, it's a good idea to support parallel transport there.

For the Cholesky-QR approach mentioned for retract_qr_fused! and Grassmann exp!: the key reference is CholeskyQR2 (Fukaya & Nakatsukasa, 2014) — run CholQR twice for stability. This is a different algorithm from Manifolds.jl's Householder QR (qr() → orgqr!), chosen because all primitives are batchable on GPU (gemm_strided_batched for A'A, potrfBatched! for Cholesky). For retraction where A ≈ I + small step, even single-pass CholQR should be well-conditioned.

Good question. Usually the QR problem on Grassmann is very well conditioned but I'm not sure it's entirely guaranteed. It would probably make most sense to introduce a new retraction method that specifies whether single-Cholesky or double-Cholesky is used.

Curiously, there is even a fancier and more stable triple Cholesky with shifting:
Shifted Cholesky QR for computing the QR factorization of ill-conditioned matrices