Test block-GMRES on GPUs

JuliaSmoothOptimizers · Sep 30, 2024 · a4da5de · a4da5de
1 parent db20870
commit a4da5de
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 1 deletion.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -11,6 +11,7 @@ steps:
         using Pkg
         Pkg.add("CUDA")
         Pkg.add("LinearOperators")
+        Pkg.add("KernelAbstractions")
         Pkg.instantiate()
         using CUDA
         # CUDA.set_runtime_version!(v"11.8")'
@@ -36,6 +37,7 @@ steps:
       julia --color=yes --project -e '
         using Pkg
         Pkg.add("AMDGPU")
+        Pkg.add("KernelAbstractions")
         Pkg.instantiate()
         include("test/gpu/amd.jl")'
     timeout_in_minutes: 30
@@ -52,6 +54,7 @@ steps:
         using Pkg
         Pkg.add(url="https://github.com/JuliaGPU/oneAPI.jl", rev="master")
         # Pkg.add("oneAPI")
+        Pkg.add("KernelAbstractions")
         Pkg.instantiate()
         include("test/gpu/intel.jl")'
     timeout_in_minutes: 30
@@ -68,6 +71,7 @@ steps:
       julia --color=yes --project -e '
         using Pkg
         Pkg.add("Metal")
+        Pkg.add("KernelAbstractions")
         Pkg.instantiate()
         include("test/gpu/metal.jl")'
     timeout_in_minutes: 30

diff --git a/test/gpu/amd.jl b/test/gpu/amd.jl
@@ -103,6 +103,15 @@ include("gpu.jl")
       @test norm(b - A * x) ≤ atol + rtol * norm(b)
     end
 
+    @testset "block-GMRES -- $FC" begin
+      A, b = nonsymmetric_indefinite(FC=FC)
+      B = hcat(b, -b)
+      A = M(A)
+      B = M(B)
+      X, stats = block_gmres(A, B)
+      @test norm(B - A * X) ≤ atol + rtol * norm(B)
+    end
+
     @testset "CG -- $FC" begin
       A, b = symmetric_definite(FC=FC)
       A = M(A)

diff --git a/test/gpu/gpu.jl b/test/gpu/gpu.jl
@@ -1,5 +1,19 @@
 using SparseArrays, Random, Test
-using LinearAlgebra, Krylov
+using LinearAlgebra, Krylov, KernelAbstractions
+
+@kernel function copy_triangle_kernel!(dest, src)
+  i, j = @index(Global, NTuple)
+  if j >= i
+    @inbounds dest[i, j] = src[i, j]
+  end
+end
+
+function Krylov.copy_triangle(Q::AbstractMatrix{FC}, R::AbstractMatrix{FC}, k::Int) where FC <: Krylov.FloatOrComplex
+  backend = get_backend(Q)
+  ndrange = (k, k)
+  copy_triangle_kernel!(backend)(R, Q; ndrange=ndrange)
+  KernelAbstractions.synchronize(backend)
+end
 
 Random.seed!(666)
 

diff --git a/test/gpu/intel.jl b/test/gpu/intel.jl
@@ -86,6 +86,15 @@ include("gpu.jl")
       @test norm(b - A * x) ≤ atol + rtol * norm(b)
     end
 
+    @testset "block-GMRES -- $FC" begin
+      A, b = nonsymmetric_indefinite(FC=FC)
+      B = hcat(b, -b)
+      A = M(A)
+      B = M(B)
+      X, stats = block_gmres(A, B)
+      @test norm(B - A * X) ≤ atol + rtol * norm(B)
+    end
+
     @testset "CG -- $FC" begin
       A, b = symmetric_definite(FC=FC)
       A = M(A)

diff --git a/test/gpu/nvidia.jl b/test/gpu/nvidia.jl
@@ -186,6 +186,15 @@ include("gpu.jl")
       @test norm(b - A * x) ≤ atol + rtol * norm(b)
     end
 
+    @testset "block-GMRES -- $FC" begin
+      A, b = nonsymmetric_indefinite(FC=FC)
+      B = hcat(b, -b)
+      A = M(A)
+      B = M(B)
+      X, stats = block_gmres(A, B)
+      @test norm(B - A * X) ≤ atol + rtol * norm(B)
+    end
+
     @testset "CG -- $FC" begin
       A, b = symmetric_definite(FC=FC)
       A = M(A)