From 200a9cd01c3bba906b82eaa87103aa2881983e5f Mon Sep 17 00:00:00 2001
From: Alexis Montoison <35051714+amontoison@users.noreply.github.com>
Date: Mon, 30 Sep 2024 02:17:37 -0500
Subject: [PATCH] Test block-GMRES on GPUs (#878)

* Test block-GMRES on GPUs

* Update intel.jl

* Update intel.jl

* Update intel.jl

* Update intel.jl
---
 .buildkite/pipeline.yml |  4 ++++
 test/gpu/amd.jl         |  9 +++++++++
 test/gpu/gpu.jl         | 16 +++++++++++++++-
 test/gpu/intel.jl       | 27 ++++++++++++++++++---------
 test/gpu/nvidia.jl      |  9 +++++++++
 5 files changed, 55 insertions(+), 10 deletions(-)

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 828eceaa5..569a59be1 100644
--- 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
index 6834d8a3f..d9b0e002d 100644
--- 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
index 1edbf9cc0..dcd7a007a 100644
--- 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
index f03176199..69da5e4fb 100644
--- a/test/gpu/intel.jl
+++ b/test/gpu/intel.jl
@@ -18,7 +18,7 @@ include("gpu.jl")
     x, stats = lsqr(A_gpu, b_gpu)
   end
 
-  for FC ∈ (Float32, ComplexF32)
+  for FC ∈ (Float32, ) # ComplexF32)
     S = oneVector{FC}
     M = oneMatrix{FC}
     T = real(FC)
@@ -79,29 +79,38 @@ include("gpu.jl")
     rtol = √ε
 
     @testset "GMRES -- $FC" begin
-      A, b = nonsymmetric_indefinite(FC=FC)
+      A, b = symmetric_definite(FC=FC)
       A = M(A)
       b = S(b)
       x, stats = gmres(A, b)
       @test norm(b - A * x) ≤ atol + rtol * norm(b)
     end
 
-    @testset "CG -- $FC" begin
+    @testset "block-GMRES -- $FC" begin
       A, b = symmetric_definite(FC=FC)
+      B = hcat(b, -b)
       A = M(A)
-      b = S(b)
-      x, stats = cg(A, b)
-      @test norm(b - A * x) ≤ atol + rtol * norm(b)
+      B = M(B)
+      X, stats = block_gmres(A, B)
+      @test norm(B - A * X) ≤ atol + rtol * norm(B)
     end
 
-    @testset "MINRES-QLP -- $FC" begin
-      A, b = symmetric_indefinite(FC=FC)
+    @testset "CG -- $FC" begin
+      A, b = symmetric_definite(FC=FC)
       A = M(A)
       b = S(b)
-      x, stats = minres_qlp(A, b)
+      x, stats = cg(A, b)
       @test norm(b - A * x) ≤ atol + rtol * norm(b)
     end
 
+    # @testset "MINRES-QLP -- $FC" begin
+    #   A, b = symmetric_indefinite(FC=FC)
+    #   A = M(A)
+    #   b = S(b)
+    #   x, stats = minres_qlp(A, b)
+    #   @test norm(b - A * x) ≤ atol + rtol * norm(b)
+    # end
+
     # @testset "processes -- $FC" begin
     #   test_processes(S, M)
     # end
diff --git a/test/gpu/nvidia.jl b/test/gpu/nvidia.jl
index a73d02062..6ae584ebf 100644
--- 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)