From 611b3c2632c40e6e77fcc4da4cbe4abc32c79246 Mon Sep 17 00:00:00 2001
From: Alexis Montoison <alexis.montoison@polymtl.ca>
Date: Fri, 23 Sep 2022 10:07:18 -0400
Subject: [PATCH] Add vector_to_matrix function

---
 src/krylov_utils.jl | 18 ++++++++++++++++++
 test/gpu/amd.jl     |  6 ++++++
 test/gpu/intel.jl   |  6 ++++++
 test/gpu/metal.jl   |  6 ++++++
 test/gpu/nvidia.jl  |  8 ++++++--
 test/test_aux.jl    |  9 +++++++++
 6 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/src/krylov_utils.jl b/src/krylov_utils.jl
index 46c9d6cd6..b16da57c0 100644
--- a/src/krylov_utils.jl
+++ b/src/krylov_utils.jl
@@ -247,6 +247,24 @@ function ktypeof(v::S) where S <: SubArray
   return ktypeof(v.parent)
 end
 
+"""
+    M = vector_to_matrix(S)
+
+Return the dense matrix storage type `M` related to the dense vector storage type `S`.
+"""
+function vector_to_matrix(::Type{S}) where S <: DenseVector
+  V = hasproperty(S, :body) ? S.body : S
+  par = V.parameters
+  npar = length(par)
+  (2 ≤ npar ≤ 3) || error("Type $S is not supported.")
+  if npar == 2
+    M = V.name.wrapper{par[1], 2}
+  else
+    M = V.name.wrapper{par[1], 2, par[3]}
+  end
+  return M
+end
+
 """
     v = kzeros(S, n)
 
diff --git a/test/gpu/amd.jl b/test/gpu/amd.jl
index 03ada1d4d..baad2bdcf 100644
--- a/test/gpu/amd.jl
+++ b/test/gpu/amd.jl
@@ -68,6 +68,12 @@ include("../test_utils.jl")
     #   Krylov.@kref!(n, x, y, c, s)
     # end
 
+    @testset "vector_to_matrix" begin
+      S = ROCVector{FC}
+      M = Krylov.vector_to_matrix(S)
+      @test M == ROCMatrix{FC}
+    end
+
     ε = eps(T)
     atol = √ε
     rtol = √ε
diff --git a/test/gpu/intel.jl b/test/gpu/intel.jl
index e6826e9e9..67ad0a7d5 100644
--- a/test/gpu/intel.jl
+++ b/test/gpu/intel.jl
@@ -76,6 +76,12 @@ end
     #   Krylov.@kref!(n, x, y, c, s)
     # end
 
+    @testset "vector_to_matrix" begin
+      S = oneVector{FC}
+      M = Krylov.vector_to_matrix(S)
+      @test M == oneMatrix{FC}
+    end
+
     ε = eps(T)
     atol = √ε
     rtol = √ε
diff --git a/test/gpu/metal.jl b/test/gpu/metal.jl
index 774ccc10c..35325c863 100644
--- a/test/gpu/metal.jl
+++ b/test/gpu/metal.jl
@@ -80,6 +80,12 @@ end
     #   Krylov.@kref!(n, x, y, c, s)
     # end
 
+    @testset "vector_to_matrix" begin
+      S = MtlVector{FC}
+      M = Krylov.vector_to_matrix(S)
+      @test M == MtlMatrix{FC}
+    end
+
     ε = eps(T)
     atol = √ε
     rtol = √ε
diff --git a/test/gpu/nvidia.jl b/test/gpu/nvidia.jl
index 8dfb61b0f..8faed479a 100644
--- a/test/gpu/nvidia.jl
+++ b/test/gpu/nvidia.jl
@@ -3,8 +3,6 @@ using LinearOperators, Krylov, CUDA, CUDA.CUSPARSE, CUDA.CUSOLVER
 
 include("../test_utils.jl")
 
-include("../test_utils.jl")
-
 @testset "Nvidia -- CUDA.jl" begin
 
   @test CUDA.functional()
@@ -146,6 +144,12 @@ include("../test_utils.jl")
       Krylov.@kref!(n, x, y, c, s)
     end
 
+    @testset "vector_to_matrix" begin
+      S = CuVector{FC}
+      M = Krylov.vector_to_matrix(S)
+      @test M == CuMatrix{FC}
+    end
+
     ε = eps(T)
     atol = √ε
     rtol = √ε
diff --git a/test/test_aux.jl b/test/test_aux.jl
index 5a4d094c7..5ac2b401c 100644
--- a/test/test_aux.jl
+++ b/test/test_aux.jl
@@ -134,6 +134,15 @@
     @test Krylov.ktypeof(b) == Vector{Float64}
   end
 
+  @testset "vector_to_matrix" begin
+    # test vector_to_matrix
+    for FC in (Float32, Float64, ComplexF32, ComplexF64)
+      S = Vector{FC}
+      M = Krylov.vector_to_matrix(S)
+      @test M == Matrix{FC}
+    end
+  end
+
   @testset "macros" begin
     # test macros
     for FC ∈ (Float16, Float32, Float64, ComplexF16, ComplexF32, ComplexF64)