add GPU project! (point + tangent) for Stiefel

Author: zazabap

ext/ManifoldsGPUCUDAExt/Stiefel.jl

@@ -48,6 +48,28 @@ function ManifoldsBase.norm(
     return sqrt(dot(X, X))
 end
 
+function ManifoldsBase.project!(
+        ::PowerManifold{ℝ, <:Stiefel{ℝ}, <:Tuple, ArrayPowerRepresentation},
+        q::CuArray{T, 3},
+        p::CuArray{T, 3},
+    ) where {T <: Real}
+    q .= p
+    return _polar_project_gpu!(q)
+end
+
+function ManifoldsBase.project!(
+        ::PowerManifold{ℝ, <:Stiefel{ℝ}, <:Tuple, ArrayPowerRepresentation},
+        Y::CuArray{T, 3},
+        p::CuArray{T, 3},
+        X::CuArray{T, 3},
+    ) where {T <: Real}
+    A = CUDA.CUBLAS.gemm_strided_batched('C', 'N', p, X)    # p'X, k×k×batch
+    sym = A .+ permutedims(A, (2, 1, 3))                     # A + A'
+    Y .= X
+    CUDA.CUBLAS.gemm_strided_batched!('N', 'N', T(-0.5), p, sym, one(T), Y)
+    return Y
+end
+
 function ManifoldsBase.retract_polar_fused!(
         ::PowerManifold{ℝ, <:Stiefel{ℝ}, <:Tuple, ArrayPowerRepresentation},
         q::CuArray{T, 3},

test/cuda/test_stiefel.jl

@@ -150,6 +150,54 @@
         @test isapprox(n_gpu, n_cpu; atol = 1.0e-10, rtol = 1.0e-10)
     end
 
+    @testset "project! point Float64" begin
+        Random.seed!(76)
+
+        M = Stiefel(8, 4)
+        MP = PowerManifold(M, 32)
+
+        p = randn(size(rand(MP))...)
+
+        q_cpu = similar(p)
+        for i in 1:size(p, 3)
+            ManifoldsBase.project!(
+                M, view(q_cpu, :, :, i), view(p, :, :, i)
+            )
+        end
+
+        p_cu = CuArray(p)
+        q_cu = similar(p_cu)
+        ManifoldsBase.project!(MP, q_cu, p_cu)
+        q_cu_h = Array(q_cu)
+
+        @test is_point(MP, q_cu_h)
+        @test isapprox(q_cu_h, q_cpu; atol = 2.0e-14, rtol = 2.0e-14)
+    end
+
+    @testset "project! point Float32" begin
+        Random.seed!(77)
+
+        M = Stiefel(8, 4)
+        MP = PowerManifold(M, 32)
+
+        p = Float32.(randn(size(rand(MP))...))
+
+        q_cpu = similar(p)
+        for i in 1:size(p, 3)
+            ManifoldsBase.project!(
+                M, view(q_cpu, :, :, i), view(p, :, :, i)
+            )
+        end
+
+        p_cu = CuArray(p)
+        q_cu = similar(p_cu)
+        ManifoldsBase.project!(MP, q_cu, p_cu)
+        q_cu_h = Array(q_cu)
+
+        @test is_point(MP, q_cu_h)
+        @test isapprox(q_cu_h, q_cpu; atol = 2.0f-5, rtol = 2.0f-5)
+    end
+
     @testset "inner and norm Float32" begin
         Random.seed!(75)
 
@@ -173,4 +221,54 @@
         @test isapprox(i_gpu, i_cpu; atol = 1.0f-4, rtol = 1.0f-4)
         @test isapprox(n_gpu, n_cpu; atol = 1.0f-4, rtol = 1.0f-4)
     end
+
+    @testset "project! tangent Float64" begin
+        Random.seed!(78)
+
+        M = Stiefel(8, 4)
+        MP = PowerManifold(M, 32)
+
+        p = rand(MP)
+        X = randn(size(p)...)
+
+        Y_cpu = similar(X)
+        for i in 1:size(p, 3)
+            ManifoldsBase.project!(
+                M, view(Y_cpu, :, :, i), view(p, :, :, i), view(X, :, :, i)
+            )
+        end
+
+        p_cu = CuArray(p)
+        X_cu = CuArray(X)
+        Y_cu = similar(X_cu)
+        ManifoldsBase.project!(MP, Y_cu, p_cu, X_cu)
+        Y_cu_h = Array(Y_cu)
+
+        @test isapprox(Y_cu_h, Y_cpu; atol = 2.0e-14, rtol = 2.0e-14)
+    end
+
+    @testset "project! tangent Float32" begin
+        Random.seed!(79)
+
+        M = Stiefel(8, 4)
+        MP = PowerManifold(M, 32)
+
+        p = Float32.(rand(MP))
+        X = Float32.(randn(size(p)...))
+
+        Y_cpu = similar(X)
+        for i in 1:size(p, 3)
+            ManifoldsBase.project!(
+                M, view(Y_cpu, :, :, i), view(p, :, :, i), view(X, :, :, i)
+            )
+        end
+
+        p_cu = CuArray(p)
+        X_cu = CuArray(X)
+        Y_cu = similar(X_cu)
+        ManifoldsBase.project!(MP, Y_cu, p_cu, X_cu)
+        Y_cu_h = Array(Y_cu)
+
+        @test isapprox(Y_cu_h, Y_cpu; atol = 2.0f-5, rtol = 2.0f-5)
+    end
 end