Add an extension for oneAPI.jl

.github/workflows/test.yml

@@ -108,7 +108,7 @@ jobs:
       fail-fast: false
       matrix:
         julia-version: ['1']
-        gpu: [cuda, amdgpu]
+        gpu: [cuda, amdgpu, oneapi]
     runs-on:
       - self-hosted
       - ${{ matrix.gpu }}
@@ -118,13 +118,30 @@ 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=".")
-          Pkg.test("MadNLPGPU", coverage=true)
+          Pkg.develop(path="./lib/MadNLPGPU")
+          Pkg.develop(path="./lib/MadNLPTests")
+      - name: Add AMDGPU
+        if: matrix.gpu == 'amdgpu'
+        shell: julia --project=./lib/MadNLPGPU/test --color=yes {0}
+        run: |
+          using Pkg
+          Pkg.add("AMDGPU")
+      - name: Add oneAPI
+        if: matrix.gpu == 'oneapi'
+        shell: julia --project=./lib/MadNLPGPU/test --color=yes {0}
+        run: |
+          using Pkg
+          Pkg.add("oneAPI")
+      - name: test MadNLPGPU
+        env:
+          MADNLP_GPU_BACKEND: ${{ matrix.gpu }}
+        run: julia --project=./lib/MadNLPGPU/test --color=yes -e 'include("lib/MadNLPGPU/test/runtests.jl")'
       - uses: julia-actions/julia-processcoverage@v1
         with:
           directories: lib/MadNLPGPU/src

lib/MadNLPGPU/Project.toml

@@ -6,6 +6,7 @@ version = "0.7.18"
 AMD = "14f7f29c-3bd6-536c-9a0b-7339e30b5a3e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e"
+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,24 +15,20 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"
 
 [extensions]
 MadNLPGPUAMDGPUExt = "AMDGPU"
+MadNLPGPUOneAPIExt = "oneAPI"
 
 [compat]
 AMD = "0.5"
 AMDGPU = "2"
 CUDA = "5.4.0"
 CUDSS = "0.6.4"
+GPUArraysCore = "0.2"
 KernelAbstractions = "0.9"
 MadNLP = "0.8.12"
-MadNLPTests = "0.5.3"
 Metis = "1"
 julia = "1.10"
-
-[extras]
-MadNLPTests = "b52a2a03-04ab-4a5f-9698-6a2deff93217"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-
-[targets]
-test = ["Test", "MadNLPTests", "AMDGPU"]
+oneAPI = "2.6.0"

lib/MadNLPGPU/ext/MadNLPGPUOneAPIExt/MadNLPGPUOneAPIExt.jl

@@ -0,0 +1,26 @@
+module MadNLPGPUOneAPIExt
+
+import LinearAlgebra
+import SparseArrays: SparseMatrixCSC, nonzeros, nnz
+import LinearAlgebra: Symmetric
+
+import MadNLP
+import MadNLPGPU
+
+import KernelAbstractions: synchronize
+import GPUArraysCore: @allowscalar
+
+using oneAPI
+using oneAPI.oneMKL, oneAPI.Support
+
+function __init__()
+    setglobal!(MadNLPGPU, :LapackOneMKLSolver, LapackOneMKLSolver)
+    return
+end
+
+include("oneapi_dense.jl")
+include("oneapi_sparse.jl")
+include("onemkl.jl")
+include("oneapi.jl")
+
+end

lib/MadNLPGPU/ext/MadNLPGPUOneAPIExt/oneapi.jl

@@ -0,0 +1,106 @@
+#=
+    MadNLP.MadNLPOptions
+=#
+
+function MadNLP.MadNLPOptions{T}(
+        nlp::MadNLP.AbstractNLPModel{T, VT};
+        dense_callback = MadNLP.is_dense_callback(nlp),
+        callback = dense_callback ? MadNLP.DenseCallback : MadNLP.SparseCallback,
+        kkt_system = dense_callback ? MadNLP.DenseCondensedKKTSystem : MadNLP.SparseCondensedKKTSystem,
+        linear_solver = MadNLPGPU.LapackOneMKLSolver,
+        tol = MadNLP.get_tolerance(T, kkt_system),
+        bound_relax_factor = tol,
+    ) where {T, VT <: oneVector{T}}
+    return MadNLP.MadNLPOptions{T}(
+        tol = tol,
+        callback = callback,
+        kkt_system = kkt_system,
+        linear_solver = linear_solver,
+        bound_relax_factor = bound_relax_factor,
+    )
+end
+
+#=
+    SparseMatrixCSC to oneSparseMatrixCSC
+=#
+
+function oneMKL.oneSparseMatrixCSC{Tv, Ti}(A::SparseMatrixCSC{Tv, Ti}) where {Tv, Ti}
+    return oneMKL.oneSparseMatrixCSC{Tv, Ti}(
+        oneVector(A.colptr),
+        oneVector(A.rowval),
+        oneVector(A.nzval),
+        size(A),
+    )
+end
+
+#=
+    oneSparseMatrixCSC to oneMatrix
+=#
+
+function MadNLPGPU.gpu_transfer!(y::oneMatrix{T}, x::oneMKL.oneSparseMatrixCSC{T}) where {T}
+    n = size(y, 2)
+    fill!(y, zero(T))
+    backend = oneAPIBackend()
+    MadNLPGPU._csc_to_dense_kernel!(backend)(y, x.colPtr, x.rowVal, x.nzVal, ndrange = n)
+    synchronize(backend)
+    return
+end
+
+#=
+    MadNLP._syr!
+=#
+
+MadNLP._syr!(uplo::Char, alpha::T, x::oneVector{T}, A::oneMatrix{T}) where {T} = oneMKL.syr!(uplo, alpha, x, A)
+
+#=
+    MadNLP._symv!
+=#
+
+MadNLP._symv!(uplo::Char, alpha::T, A::oneMatrix{T}, x::oneVector{T}, beta::T, y::oneVector{T}) where {T} = oneMKL.symv!(uplo, alpha, A, x, beta, y)
+
+#=
+    MadNLP._syrk!
+=#
+
+MadNLP._syrk!(uplo::Char, trans::Char, alpha::T, A::oneMatrix{T}, beta::T, C::oneMatrix{T}) where {T} = oneMKL.syrk!(uplo, trans, alpha, A, beta, C)
+
+#=
+    MadNLP._trsm!
+=#
+
+MadNLP._trsm!(side::Char, uplo::Char, transa::Char, diag::Char, alpha::T, A::oneMatrix{T}, B::oneMatrix{T}) where {T} = oneMKL.trsm!(side, uplo, transa, diag, alpha, A, B)
+
+#=
+    MadNLP._dgmm!
+=#
+
+MadNLP._dgmm!(side::Char, A::oneMatrix{T}, x::oneVector{T}, B::oneMatrix{T}) where {T} = oneMKL.dgmm!(side, A, x, B)
+
+#=
+    LinearAlgebra.norm for oneAPI arrays.
+    GPUArrays' norm uses LinearAlgebra.norm as a map function in mapreduce,
+    which fails on oneAPI because norm on scalars generates jl_f_throw_methoderror
+    calls that can't be compiled to SPIRV. Replace with abs-based implementation.
+    The SubArray method also avoids scalar indexing on views.
+=#
+
+function _onenorm(v, p::Real)
+    isempty(v) && return float(zero(eltype(v)))
+    if p == Inf
+        return maximum(abs, v)
+    elseif p == -Inf
+        return minimum(abs, v)
+    elseif p == 1
+        return mapreduce(abs, +, v)
+    elseif p == 2
+        return sqrt(mapreduce(abs2, +, v))
+    elseif p == 0
+        return float(count(!iszero, v))
+    else
+        spp = float(p)
+        return mapreduce(x -> abs(x)^spp, +, v)^inv(spp)
+    end
+end
+
+LinearAlgebra.norm(v::oneArray{<:Number}, p::Real = 2) = _onenorm(v, p)
+LinearAlgebra.norm(v::SubArray{<:Number, <:Any, <:oneArray}, p::Real = 2) = _onenorm(v, p)

lib/MadNLPGPU/ext/MadNLPGPUOneAPIExt/oneapi_dense.jl

@@ -0,0 +1,149 @@
+########################################################################
+##### oneAPI wrappers for DenseKKTSystem / DenseCondensedKKTSystem #####
+########################################################################
+
+#=
+    MadNLP._ger!
+=#
+
+MadNLP._ger!(alpha::T, x::oneVector{T}, y::oneVector{T}, A::oneMatrix{T}) where {T} = oneMKL.ger!(alpha, x, y, A)
+
+#=
+    MadNLP._madnlp_unsafe_wrap
+=#
+
+function MadNLP._madnlp_unsafe_wrap(vec::VT, n, shift = 1) where {T, VT <: oneVector{T}}
+    return view(vec, shift:(shift + n - 1))
+end
+
+#=
+    MadNLP.diag!
+=#
+
+function MadNLP.diag!(dest::oneVector{T}, src::oneMatrix{T}) where {T}
+    @assert length(dest) == size(src, 1)
+    backend = oneAPIBackend()
+    MadNLPGPU._copy_diag_kernel!(backend)(dest, src, ndrange = length(dest))
+    synchronize(backend)
+    return
+end
+
+#=
+    MadNLP.diag_add!
+=#
+
+function MadNLP.diag_add!(dest::oneMatrix, src1::oneVector, src2::oneVector)
+    backend = oneAPIBackend()
+    MadNLPGPU._add_diagonal_kernel!(backend)(dest, src1, src2, ndrange = size(dest, 1))
+    synchronize(backend)
+    return
+end
+
+#=
+    MadNLP._set_diag!
+=#
+
+function MadNLP._set_diag!(A::oneMatrix, inds, a)
+    if !isempty(inds)
+        backend = oneAPIBackend()
+        MadNLPGPU._set_diag_kernel!(backend)(A, inds, a; ndrange = length(inds))
+        synchronize(backend)
+    end
+    return
+end
+
+#=
+    MadNLP._build_dense_kkt_system!
+=#
+
+function MadNLP._build_dense_kkt_system!(
+        dest::oneMatrix,
+        hess::oneMatrix,
+        jac::oneMatrix,
+        pr_diag::oneVector,
+        du_diag::oneVector,
+        diag_hess::oneVector,
+        ind_ineq::AbstractVector,
+        n,
+        m,
+        ns,
+    )
+    ind_ineq_gpu = oneVector(ind_ineq)
+    ndrange = (n + m + ns, n)
+    backend = oneAPIBackend()
+    MadNLPGPU._build_dense_kkt_system_kernel!(backend)(
+        dest,
+        hess,
+        jac,
+        pr_diag,
+        du_diag,
+        diag_hess,
+        ind_ineq_gpu,
+        n,
+        m,
+        ns,
+        ndrange = ndrange,
+    )
+    synchronize(backend)
+    return
+end
+
+#=
+    MadNLP._build_ineq_jac!
+=#
+
+function MadNLP._build_ineq_jac!(
+        dest::oneMatrix,
+        jac::oneMatrix,
+        diag_buffer::oneVector,
+        ind_ineq::AbstractVector,
+        n,
+        m_ineq,
+    )
+    (m_ineq == 0) && return # nothing to do if no ineq. constraints
+    ind_ineq_gpu = oneVector(ind_ineq)
+    ndrange = (m_ineq, n)
+    backend = oneAPIBackend()
+    MadNLPGPU._build_jacobian_condensed_kernel!(backend)(
+        dest,
+        jac,
+        diag_buffer,
+        ind_ineq_gpu,
+        m_ineq,
+        ndrange = ndrange,
+    )
+    synchronize(backend)
+    return
+end
+
+#=
+    MadNLP._build_condensed_kkt_system!
+=#
+
+function MadNLP._build_condensed_kkt_system!(
+        dest::oneMatrix,
+        hess::oneMatrix,
+        jac::oneMatrix,
+        pr_diag::oneVector,
+        du_diag::oneVector,
+        ind_eq::AbstractVector,
+        n,
+        m_eq,
+    )
+    ind_eq_gpu = oneVector(ind_eq)
+    ndrange = (n + m_eq, n)
+    backend = oneAPIBackend()
+    MadNLPGPU._build_condensed_kkt_system_kernel!(backend)(
+        dest,
+        hess,
+        jac,
+        pr_diag,
+        du_diag,
+        ind_eq_gpu,
+        n,
+        m_eq,
+        ndrange = ndrange,
+    )
+    synchronize(backend)
+    return
+end