Move CUDA/CUDSS to weak deps and deduplicate GPU backend code

.github/workflows/test.yml

@@ -119,16 +119,31 @@ jobs:
         with:
           channel: ${{ matrix.julia-version }}
       - uses: julia-actions/cache@v1
-      - name: test MadNLPGPU
-        shell: julia --project=./lib/MadNLPGPU --color=yes {0}
+      - name: Setup test environment
+        shell: julia --project=./lib/MadNLPGPU/test --color=yes {0}
         run: |
           using Pkg
           Pkg.Registry.update()
           Pkg.develop(path=".")
-          using MadNLP
-          PATH_MADNLP = dirname(pathof(MadNLP))
-          Pkg.develop(path=joinpath(PATH_MADNLP, "..", "lib", "MadNLPTests"))
-          Pkg.test("MadNLPGPU", coverage=true)
+          Pkg.develop(path="./lib/MadNLPGPU")
+          Pkg.develop(path="./lib/MadNLPTests")
+      - name: Add CUDA and CUDSS
+        if: matrix.gpu == 'cuda'
+        shell: julia --project=./lib/MadNLPGPU/test --color=yes {0}
+        run: |
+          using Pkg
+          Pkg.add("CUDA")
+          Pkg.add("CUDSS")
+      - name: Add AMDGPU
+        if: matrix.gpu == 'amdgpu'
+        shell: julia --project=./lib/MadNLPGPU/test --color=yes {0}
+        run: |
+          using Pkg
+          Pkg.add("AMDGPU")
+      - name: test MadNLPGPU
+        env:
+          MADNLP_GPU_BACKEND: ${{ matrix.gpu }}
+        run: julia --project=./lib/MadNLPGPU/test --code-coverage=user --color=yes -e 'include("lib/MadNLPGPU/test/runtests.jl")'
       - uses: julia-actions/julia-processcoverage@v1
         with:
           directories: lib/MadNLPGPU/src

docs/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ExaModels = "1037b233-b668-4ce9-9b63-f9f681f55dd2"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"

docs/src/tutorials/gpu.md

@@ -15,6 +15,7 @@ by importing MadNLPGPU:
 ```@example gpu
 using MadNLPGPU
 using CUDA
+using CUDSS
 
 ```
 

lib/MadNLPGPU/Project.toml

@@ -4,8 +4,8 @@ version = "0.8.0"
 
 [deps]
 AMD = "14f7f29c-3bd6-536c-9a0b-7339e30b5a3e"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e"
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6"
@@ -14,15 +14,20 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e"
 
 [extensions]
 MadNLPGPUAMDGPUExt = "AMDGPU"
+MadNLPGPUCUDAExt = ["CUDA", "CUDSS"]
 
 [compat]
 AMD = "0.5"
+Adapt = "4"
 AMDGPU = "2"
 CUDA = "5.4.0"
 CUDSS = "0.6.4"
+GPUArraysCore = "0.2"
 KernelAbstractions = "0.9"
 MadNLP = "0.9"
 MadNLPTests = "0.6"
@@ -34,4 +39,4 @@ MadNLPTests = "b52a2a03-04ab-4a5f-9698-6a2deff93217"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "MadNLPTests", "AMDGPU"]
+test = ["Test", "MadNLPTests"]

lib/MadNLPGPU/ext/MadNLPGPUAMDGPUExt/MadNLPGPUAMDGPUExt.jl

@@ -7,7 +7,7 @@ import LinearAlgebra: Symmetric
 import MadNLP
 import MadNLPGPU
 
-import KernelAbstractions: synchronize
+import KernelAbstractions: synchronize, get_backend
 
 using AMDGPU
 using AMDGPU.rocBLAS, AMDGPU.rocSOLVER, AMDGPU.rocSPARSE
@@ -17,9 +17,7 @@ function __init__()
     return
 end
 
-include("rocm_dense.jl")
 include("rocm_sparse.jl")
-include("rocm_qn.jl")
 include("rocsolver.jl")
 include("rocm.jl")
 

lib/MadNLPGPU/ext/MadNLPGPUAMDGPUExt/rocm.jl

@@ -37,29 +37,8 @@ end
     ROCSparseMatrixCSC to ROCMatrix
 =#
 
-function MadNLPGPU.gpu_transfer!(y::ROCMatrix{T}, x::rocSPARSE.ROCSparseMatrixCSC{T}) where {T}
-    n = size(y, 2)
-    fill!(y, zero(T))
-    backend = ROCBackend()
-    MadNLPGPU._csc_to_dense_kernel!(backend)(y, x.colPtr, x.rowVal, x.nzVal, ndrange = n)
-    return
-end
-
-if VERSION > v"1.11" # See https://github.com/JuliaGPU/AMDGPU.jl/issues/607. 1-norm of view() of RocArray seems to have been broken in v1.12
-    function MadNLP.get_sd(l::ROCVector{T}, zl_r, zu_r, s_max) where T
-        return max(
-            s_max,
-            (my1norm(l)+my1norm(zl_r)+my1norm(zu_r)) / max(1, (length(l)+length(zl_r)+length(zu_r))),
-        ) / s_max
-    end
-    function MadNLP.get_sc(zl_r::SubArray{T,1,VT}, zu_r, s_max) where {T, VT <: ROCVector{T}}
-        return max(
-            s_max,
-            (my1norm(zl_r)+my1norm(zu_r)) / max(1,length(zl_r)+length(zu_r)),
-        ) / s_max
-    end
-    my1norm(x) = mapreduce(abs, +, x)
-end
+MadNLPGPU.gpu_transfer!(y::ROCMatrix{T}, x::rocSPARSE.ROCSparseMatrixCSC{T}) where {T} =
+    MadNLPGPU._gpu_sparse_csc_to_dense!(y, x)
 
 #=
     MadNLP._ger!