GPU-native manifold operations via CUDA extension

Project.toml

@@ -24,6 +24,7 @@ Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 
 [weakdeps]
 BoundaryValueDiffEq = "764a87c0-6b3e-53db-9096-fe964310641d"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -36,6 +37,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [extensions]
 ManifoldsBoundaryValueDiffEqExt = "BoundaryValueDiffEq"
+ManifoldsCUDAExt = "CUDA"
 ManifoldsDistributionsExt = ["Distributions", "RecursiveArrayTools"]
 ManifoldsHybridArraysExt = "HybridArrays"
 ManifoldsNLsolveExt = "NLsolve"
@@ -47,6 +49,7 @@ ManifoldsTestExt = "Test"
 
 [compat]
 ADTypes = "1.19.0"
+CUDA = "5"
 BoundaryValueDiffEq = "4, 5.6.1"
 Colors = "0.12, 0.13"
 DiffEqCallbacks = "2, 3, 4"

ext/ManifoldsCUDAExt.jl

@@ -0,0 +1,136 @@
+"""
+    ManifoldsCUDAExt
+
+CUDA extension for Manifolds.jl providing GPU-native overrides for:
+- `exp!` on `GeneralUnitaryMatrices` via Hermitian eigendecomposition
+- `log_safe!` via eigendecomposition on GPU (cuSOLVER `geev!`/`heevd!`)
+- QR retraction on `GeneralUnitaryMatrices` and `Grassmann`
+- Random point generation on `UnitaryMatrices`
+"""
+module ManifoldsCUDAExt
+
+using Manifolds
+using ManifoldsBase
+using ManifoldsBase: AbstractManifold
+import ManifoldsBase: retract_qr_fused!
+
+using Manifolds:
+    GeneralUnitaryMatrices,
+    UnitaryMatrices,
+    Grassmann
+
+using CUDA
+using LinearAlgebra
+using Random
+
+# GPU-native matrix exponential for skew-Hermitian tangent vectors.
+#
+# For skew-Hermitian X (X' = -X), the matrix iX is Hermitian.
+# Using eigen(Hermitian(iX)) = (V, λ) where λ are real:
+#   exp(X) = V * Diagonal(exp(-i*λ)) * V'
+#
+# For real skew-symmetric X, we promote to complex, compute, and take real part.
+function Manifolds.exp!(
+    M::GeneralUnitaryMatrices, q::CuArray{T}, p::CuArray{T}, X::CuArray{T}
+) where {T<:Complex}
+    iX = Hermitian(im .* X)
+    F = eigen(iX)
+    expX = F.vectors * Diagonal(exp.(-im .* F.values)) * F.vectors'
+    q .= p * expX
+    return q
+end
+
+function Manifolds.exp!(
+    M::GeneralUnitaryMatrices, q::CuArray{T}, p::CuArray{T}, X::CuArray{T}
+) where {T<:Real}
+    Xc = CuArray{complex(T)}(X)
+    iX = Hermitian(im .* Xc)
+    F = eigen(iX)
+    expX = F.vectors * Diagonal(exp.(-im .* F.values)) * F.vectors'
+    q .= real.(CuArray{complex(T)}(p) * expX)
+    return q
+end
+
+# GPU-native matrix logarithm via eigendecomposition.
+#
+# cuSOLVER provides geev! for general eigendecomposition on GPU.
+# For a matrix A with eigendecomposition A = V * Diagonal(λ) * V⁻¹:
+#   log(A) = V * Diagonal(log(λ)) * V⁻¹
+function Manifolds.log_safe!(Y::CuArray{T}, A::CuArray{T}) where {T<:Complex}
+    F = eigen(A)
+    Y .= F.vectors * Diagonal(log.(F.values)) * inv(F.vectors)
+    return Y
+end
+
+function Manifolds.log_safe!(Y::CuArray{T}, A::CuArray{T}) where {T<:Real}
+    Ac = CuArray{complex(T)}(A)
+    F = eigen(Ac)
+    logA = F.vectors * Diagonal(log.(F.values)) * inv(F.vectors)
+    Y .= real.(logA)
+    return Y
+end
+
+# GPU-native QR retraction for GeneralUnitaryMatrices.
+# Base uses Diagonal(sign.(diag(R))) which triggers scalar indexing.
+function ManifoldsBase.retract_qr_fused!(
+    ::GeneralUnitaryMatrices,
+    q::CuArray{T},
+    p::CuArray{T},
+    X::CuArray{T},
+    t::Number,
+) where {T}
+    A = p + p * (T(t) * X)
+    qr_decomp = qr(A)
+    Q_mat = CuArray(qr_decomp.Q)
+    R_mat = CuArray(qr_decomp.R)
+    d = diag(R_mat)
+    signs = sign.(d .+ T(0.5))
+    q .= Q_mat .* reshape(signs, 1, :)
+    return q
+end
+
+# GPU-native QR retraction for Grassmann.
+# Base uses Array(qr.Q) creating a CPU matrix in GPU broadcast.
+function ManifoldsBase.retract_qr_fused!(
+    ::Grassmann,
+    q::CuArray{T},
+    p::CuArray{T},
+    X::CuArray{T},
+    t::Number,
+) where {T}
+    q .= p .+ T(t) .* X
+    qr_decomp = qr(q)
+    Q_mat = CuArray(qr_decomp.Q)
+    R_mat = CuArray(qr_decomp.R)
+    d = diag(R_mat)
+    signs = sign.(d .+ T(0.5))
+    q .= Q_mat .* reshape(signs, 1, :)
+    return q
+end
+
+# GPU-native random point generation for UnitaryMatrices.
+# Base rand! uses randn(rng, T, n, n) which creates CPU arrays.
+function Random.rand!(
+    rng::AbstractRNG,
+    M::UnitaryMatrices,
+    pX::CuArray;
+    vector_at = nothing,
+    σ::Real = one(real(eltype(pX))),
+)
+    n = ManifoldsBase.get_parameter(M.size)[1]
+    if vector_at === nothing
+        A = CUDA.randn(eltype(pX), n, n) .* σ
+        qr_decomp = qr(A)
+        Q_mat = CuArray(qr_decomp.Q)
+        R_mat = CuArray(qr_decomp.R)
+        d = diag(R_mat)
+        signs = sign.(d)
+        pX .= Q_mat .* reshape(signs, 1, :)
+    else
+        Z = CUDA.randn(eltype(pX), size(pX)...) .* σ
+        Manifolds.project!(M, pX, vector_at, Z)
+    end
+    return pX
+end
+
+end # module

test/runtests.jl

@@ -88,4 +88,5 @@ end
         include_test("manifolds-old/vector_bundle.jl")
         include_test("manifolds-old/graph.jl")
     end
+    include_test("test_cuda_ext.jl")
 end