From a10a5e3dd20638889dec8a1ca0cef7aae0953a9e Mon Sep 17 00:00:00 2001
From: Alexis Montoison <alexis.montoison@polymtl.ca>
Date: Tue, 28 Nov 2023 11:57:37 -0600
Subject: [PATCH] Add more tests with KrylovOperators

---
 test/gpu/gpu.jl    | 27 ++++++++++++++++++++++++---
 test/gpu/nvidia.jl |  6 +++---
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/test/gpu/gpu.jl b/test/gpu/gpu.jl
index 65453d5..4f345d4 100644
--- a/test/gpu/gpu.jl
+++ b/test/gpu/gpu.jl
@@ -42,14 +42,14 @@ function test_ilu0(FC, V, M)
   @test norm(r_gpu) ≤ 1e-8
 end
 
-function test_operator(FC, V, M)
+function test_operator(FC, V, DM, SM)
+  m = 200
   n = 100
-  R = real(FC)
   A_cpu = rand(FC, n, n)
   A_cpu = sparse(A_cpu)
   b_cpu = rand(FC, n)
 
-  A_gpu = M(A_cpu)
+  A_gpu = SM(A_cpu)
   b_gpu = V(b_cpu)
   opA_gpu = KrylovOperator(A_gpu)
 
@@ -57,6 +57,27 @@ function test_operator(FC, V, M)
   r_gpu = b_gpu - A_gpu * x_gpu
   @test stats.solved
   @test norm(r_gpu) ≤ 1e-8
+
+  A_cpu = rand(FC, m, n)
+  A_cpu = sparse(A_cpu)
+  opA_gpu = KrylovOperator(A_gpu)
+  for i = 1:5
+    y_cpu = rand(FC, m)
+    y_gpu = V(y_cpu)
+    x_cpu = rand(FC, n)
+    x_gpu = V(x_cpu)
+    mul!(y_gpu, opA_gpu, x_gpu)
+  end
+
+  nrhs = 3
+  opA_gpu = KrylovOperator(A_gpu; nrhs)
+  for i = 1:5
+    Y_cpu = rand(FC, m, nrhs)
+    Y_gpu = DM(Y_cpu)
+    X_cpu = rand(FC, n, nrhs)
+    X_gpu = DM(X_cpu)
+    mul!(Y_gpu, opA_gpu, X_gpu)
+  end
 end
 
 _get_type(J::SparseMatrixCSC) = Vector{Float64}
diff --git a/test/gpu/nvidia.jl b/test/gpu/nvidia.jl
index dbf212f..655e1df 100644
--- a/test/gpu/nvidia.jl
+++ b/test/gpu/nvidia.jl
@@ -28,13 +28,13 @@ include("gpu.jl")
 
   @testset "KrylovOperator" begin
     @testset "CuSparseMatrixCOO -- $FC" for FC in (Float64, ComplexF64)
-      test_operator(FC, CuVector{FC}, CuSparseMatrixCOO{FC})
+      test_operator(FC, CuVector{FC}, CuMatrix{FC}, CuSparseMatrixCOO{FC})
     end
     @testset "CuSparseMatrixCSC -- $FC" for FC in (Float64, ComplexF64)
-      test_operator(FC, CuVector{FC}, CuSparseMatrixCSC{FC})
+      test_operator(FC, CuVector{FC}, CuMatrix{FC}, CuSparseMatrixCSC{FC})
     end
     @testset "CuSparseMatrixCSR -- $FC" for FC in (Float64, ComplexF64)
-      test_operator(FC, CuVector{FC}, CuSparseMatrixCSR{FC})
+      test_operator(FC, CuVector{FC}, CuMatrix{FC}, CuSparseMatrixCSR{FC})
     end
   end