From 460cf12a6a546aeedf87ae33fa168b5be0737fb6 Mon Sep 17 00:00:00 2001
From: tmigot <tangi.migot@gmail.com>
Date: Tue, 30 Apr 2024 00:09:55 +0200
Subject: [PATCH] Add GPU tests

---
 Project.toml       |  3 ++-
 src/slack-model.jl |  6 +++---
 test/gpu_test.jl   | 42 ++++++++++++++++++++++++++++++++++++++++++
 test/runtests.jl   |  6 +++++-
 4 files changed, 52 insertions(+), 5 deletions(-)
 create mode 100644 test/gpu_test.jl

diff --git a/Project.toml b/Project.toml
index c949572..40052af 100644
--- a/Project.toml
+++ b/Project.toml
@@ -17,8 +17,9 @@ NLPModels = "0.19, 0.20, 0.21"
 julia = "^1.6.0"
 
 [extras]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 NLPModelsTest = "7998695d-6960-4d3a-85c4-e1bceb8cd856"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["NLPModelsTest", "Test"]
+test = ["CUDA", "NLPModelsTest", "Test"]
diff --git a/src/slack-model.jl b/src/slack-model.jl
index 2313c71..f422cb8 100644
--- a/src/slack-model.jl
+++ b/src/slack-model.jl
@@ -67,7 +67,7 @@ The slack variables are implicitly ordered as linear and then nonlinear, and
 ``c_L ≤ c(x) < ∞``, ``-∞ < c(x) ≤ c_U`` and
 ``c_L ≤ c(x) ≤ c_U``, respectively.
 """
-mutable struct SlackModel{T, S, M <: AbstractNLPModel{T, S}} <: AbstractNLPModel{T, S}
+mutable struct SlackModel{T, S, M <: AbstractNLPModel{T}} <: AbstractNLPModel{T, S}
   meta::NLPModelMeta{T, S}
   model::M
 
@@ -90,7 +90,7 @@ end
 
 """Like `SlackModel`, this model converts inequalities into equalities and bounds.
 """
-mutable struct SlackNLSModel{T, S, M <: AbstractNLPModel{T, S}} <: AbstractNLSModel{T, S}
+mutable struct SlackNLSModel{T, S, M <: AbstractNLPModel{T}} <: AbstractNLSModel{T, S}
   meta::NLPModelMeta{T, S}
   nls_meta::NLSMeta{T, S}
   model::M
@@ -177,7 +177,7 @@ function SlackNLSModel(
   x0 = similar(model.meta.x0, model.meta.nvar + ns)
   x0[(model.meta.nvar + 1):end] .= zero(T)
   x0[1:(model.meta.nvar)] .= model.meta.x0
-  nls_meta = NLSMeta{T, S}(
+  nls_meta = NLSMeta{T, typeof(meta.x0)}( # might differ from S
     model.nls_meta.nequ,
     model.meta.nvar + ns,
     x0 = x0,
diff --git a/test/gpu_test.jl b/test/gpu_test.jl
new file mode 100644
index 0000000..c87bf53
--- /dev/null
+++ b/test/gpu_test.jl
@@ -0,0 +1,42 @@
+function nlp_gpu_tests(p, Model; exclude = [])
+  @testset "NLP tests of problem $p" begin
+    nlp_from_T = T -> Model(eval(Symbol(p))(T))
+    @testset "GPU multiple precision support of problem $p" begin
+      CUDA.allowscalar() do
+        multiple_precision_nlp_array(nlp_from_T, CuArray, linear_api = true, exclude = exclude)
+      end
+    end
+  end
+end
+
+function nls_gpu_tests(p, Model; exclude = [])
+  @testset "NLP tests of problem $p" begin
+    nls_from_T = T -> Model(eval(Symbol(p))(T))
+    exclude = p == "LLS" ? union(exclude, [hess_coord, hess, hess_residual]) : exclude
+    @testset "GPU multiple precision support of problem $p" begin
+      CUDA.allowscalar() do
+        multiple_precision_nls_array(nls_from_T, CuArray, linear_api = true, exclude = exclude)
+      end
+    end
+  end
+end
+
+@testset "Check GPU multiprecision for quasi-Newton model modifiers $M of NLP" for M in [LBFGSModel, LSR1Model, DiagonalPSBModel, DiagonalAndreiModel, SpectralGradientModel]
+  map(p -> nlp_gpu_tests(p, M, exclude = [hess, hess_coord, ghjvprod, jth_hess, jth_hess_coord, jth_hprod]), union(NLPModelsTest.nlp_problems, NLPModelsTest.nls_problems))
+end
+
+@testset "Check GPU multiprecision for model modifiers $M of NLP" for M in [FeasibilityResidual]
+  map(p -> nlp_gpu_tests(p, M, exclude = [hess, hess_coord, ghjvprod, jth_hess, jth_hess_coord, jth_hprod]), setdiff(NLPModelsTest.nlp_problems, ["BROWNDEN", "HS5"]))
+end
+
+@testset "Check GPU multiprecision for model modifiers $M of NLP" for M in [SlackModel]
+  map(p -> nlp_gpu_tests(p, M, exclude = []), NLPModelsTest.nlp_problems)
+end
+
+@testset "Check GPU multiprecision for model modifiers $M of NLP" for M in [FeasibilityResidual]
+  map(p -> nls_gpu_tests(p, M, exclude = [hess, hess_coord, ghjvprod, jth_hess, jth_hess_coord, jth_hprod]), setdiff(NLPModelsTest.nls_problems, ["MGH01", "BNDROSENBROCK"]))
+end
+
+@testset "Check GPU multiprecision for model modifiers $M of NLP" for M in [SlackNLSModel, FeasibilityFormNLS]
+  map(p -> nls_gpu_tests(p, M, exclude = []), NLPModelsTest.nls_problems)
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 7aedb9b..281269e 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -10,7 +10,11 @@ include("nls/feasibility-form-nls.jl")
 include("nls/feasibility-residual.jl")
 include("nls/slack-model.jl")
 
-using NLPModelsTest
+using CUDA, NLPModelsTest
+
+if CUDA.functional()
+  include("gpu_test.jl")
+end
 
 if (v"1.7" <= VERSION)
   include("allocs_test.jl")