[Dev] AMD GPU support

[jbcaillau]

WIP MadNLP/MadNLP.jl#380

[jbcaillau]

use MadNLP + LBFGS with suitable linear solver on AMD (@amontoison)
https://madsuite.org/MadNLP.jl/dev/tutorials/lbfgs

[amontoison]

@jbcaillau

# goddard.jl

using OptimalControl
using MadNLP
using MadNLPGPU
using AMDGPU
using BenchmarkTools
using Interpolations

const r0 = 1.0
const v0 = 0.0
const m0 = 1.0
const vmax = 0.1
const mf = 0.6
const Cd = 310.0
const Tmax = 3.5
const β = 500.0
const b = 2.0

parsing_backend!(:exa)

o = @def begin

    tf ∈ R, variable
    t ∈ [0, tf], time
    x = (r, v, m) ∈ R³, state
    u ∈ R, control

    x(0) == [r0, v0, m0]
    m(tf) == mf
    0 ≤ u(t) ≤ 1
    r(t) ≥ r0
    0 ≤ v(t) ≤ vmax

    ∂(r)(t) == v(t)
    ∂(v)(t) == -Cd * v(t)^2 * exp(-β * (r(t) - 1)) / m(t) - 1 / r(t)^2 + u(t) * Tmax / m(t)
    ∂(m)(t) == -b * Tmax * u(t)

    r(tf) → max

end

tfs = 0.18761155665063417
xs0 = [ 1.0          1.00105   1.00398   1.00751    1.01009    1.01124
       -1.83989e-40  0.056163  0.1       0.0880311  0.0492518  0.0123601
        1.0          0.811509  0.650867  0.6        0.6        0.6 ]
us0 = [0.599377 0.835887 0.387328 -5.87733e-9 -9.03538e-9 -8.62101e-9]
N0 = length(us0) - 1
_t = tfs * 0:N0
_xs = linear_interpolation(_t, [xs0[:, j] for j ∈ 1:N0+1], extrapolation_bc=Line())
_us = linear_interpolation(_t, [us0[:, j] for j ∈ 1:N0+1], extrapolation_bc=Line())

tol = 1e-7
print_level = MadNLP.WARN)

N = 100

t = tfs * 0:N
xs = _xs.(t); xs = stack(xs[:])
us = _us.(t); us = stack(us[:])
m_cpu = o(; grid_size = N, init = (tfs, xs, us))
m_gpu = o(; grid_size = N, init = (tfs, xs, us), backend = ROCBackend())

printstyled("\nsolver = madnlp, N = ", N, "\n"; bold = true)

print("CPU:")
madnlp(m_cpu, linear_solver=MadNLP.LapackCPUSolver)

AMDGPU.functional() || throw("AMDGPU not available")
print("GPU:")
madnlp(m_gpu, linear_solver=MadNLPGPU.LapackROCSolver)

This is MadNLP version v0.8.9, running with rocSOLVER v3.28.0 -- (LU)

Number of nonzeros in constraint Jacobian............:     1904
Number of nonzeros in Lagrangian Hessian.............:     6100

Total number of variables............................:      405
                     variables with only lower bounds:      101
                variables with lower and upper bounds:      202
                     variables with only upper bounds:        0
Total number of equality constraints.................:      304
Total number of inequality constraints...............:        0
        inequality constraints with only lower bounds:        0
   inequality constraints with lower and upper bounds:        0
        inequality constraints with only upper bounds:        0

iter    objective    inf_pr   inf_du inf_compl lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls
   0 -1.1204900e+00 1.79e-01 2.00e+00 9.90e-01  -1.0 0.00e+00    -  0.00e+00 0.00e+00   0
   1 -1.1170310e+00 1.79e-01 1.99e+00 9.90e-01  -1.0 1.31e+01    -  3.46e-04 4.57e-04h  1
   2 -1.1148839e+00 1.79e-01 2.08e+00 9.87e-01  -1.0 8.88e+00    -  3.46e-03 8.66e-04h  1
   3 -1.1045307e+00 1.77e-01 4.95e+01 9.35e-01  -1.0 2.49e+00    -  6.48e-02 1.06e-02h  1
   4 -1.0963656e+00 1.61e-01 5.78e+01 8.94e-01  -1.0 2.56e+00    -  4.48e-02 8.61e-02h  1
   5 -1.0822020e+00 1.34e-01 1.05e+02 8.30e-01  -1.0 2.32e+00    -  6.14e-02 1.70e-01h  1
   6 -1.0716218e+00 1.13e-01 1.07e+02 3.30e-01  -1.0 5.02e-01    -  5.49e-01 1.58e-01h  1
   7 -1.0158781e+00 1.55e-03 7.61e+01 1.04e-01  -1.0 3.78e-01    -  6.07e-01 1.00e+00h  1
   8 -1.0167419e+00 3.54e-04 2.37e+02 6.97e-02  -1.0 2.33e-01    -  4.43e-01 7.54e-01h  1
   9 -1.0167178e+00 2.80e-04 1.79e+00 5.64e-02  -1.0 1.24e-01    -  1.00e+00 1.00e+00h  1
iter    objective    inf_pr   inf_du inf_compl lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls
  10 -1.0163157e+00 1.04e-05 9.69e-02 3.72e-02  -1.0 8.92e-02    -  1.00e+00 1.00e+00h  1
  11 -1.0163259e+00 4.30e-10 1.40e-02 3.25e-03  -2.5 1.22e-03    -  1.00e+00 1.00e+00h  1
  12 -1.0168085e+00 8.14e-07 1.22e-02 2.85e-04  -3.8 1.43e-02    -  1.00e+00 1.00e+00h  1
  13 -1.0246345e+00 2.53e-04 3.41e-03 1.85e-04  -3.8 2.20e-01    -  1.00e+00 1.00e+00h  1
  14 -1.0238358e+00 3.85e-06 3.62e-05 1.52e-04  -3.8 4.54e-02    -  1.00e+00 1.00e+00h  1
  15 -1.0290986e+00 1.03e-04 6.25e-01 5.02e-05  -5.0 1.23e-01    -  8.25e-01 1.00e+00h  1
  16 -1.0359760e+00 2.25e-04 4.28e-02 1.83e-05  -5.0 1.45e-01    -  1.00e+00 8.87e-01h  1
  17 -1.0385028e+00 3.65e-04 1.09e-02 1.07e-05  -5.0 2.93e-01    -  1.00e+00 9.51e-01h  1
  18 -1.0383421e+00 1.81e-05 1.03e-04 1.07e-05  -5.0 1.37e-01    -  1.00e+00 1.00e+00h  1
  19 -1.0383417e+00 1.75e-07 7.91e-07 9.14e-06  -5.0 3.30e-02    -  1.00e+00 1.00e+00h  1

Number of Iterations....: 19

                                   (scaled)                 (unscaled)
Objective...............:  -1.0383416743321257e+00   -1.0383416743321257e+00
Dual infeasibility......:   7.9144150480080395e-07    7.9144150480080395e-07
Constraint violation....:   1.7485807277055451e-07    1.7485807277055451e-07
Complementarity.........:   9.1437383333437779e-06    9.1437383333437779e-06
Overall NLP error.......:   9.1437383333437779e-06    9.1437383333437779e-06

Number of objective function evaluations             = 20
Number of objective gradient evaluations             = 20
Number of constraint evaluations                     = 20
Number of constraint Jacobian evaluations            = 20
Number of Lagrangian Hessian evaluations             = 19
Total wall-clock secs in solver (w/o fun. eval./lin. alg.)  =  0.109
Total wall-clock secs in linear solver                      =  0.181
Total wall-clock secs in NLP function evaluations           =  0.016
Total wall-clock secs                                       =  0.306

EXIT: Optimal Solution Found (tol = 1.0e-04).
"Execution stats: Optimal Solution Found (tol = 1.0e-04)."

[jbcaillau]

@amontoison awesome!

@michel2323 would ROC(K) to add an AMD runner somewhere 🤞🏽

[jbcaillau]

•  @jbcaillau test and update doc

[michel2323]

@michel2323 would ROC(K) to add an AMD runner somewhere 🤞🏽

Added. It's called pinwheel and has the amdgpu label. https://github.com/organizations/control-toolbox/settings/actions/runners

[michel2323]

#484

[michel2323]

@amontoison What is missing for Intel GPUs?

[amontoison]

@michel2323
For ExaModels, it should work without any issue.
It only relies on KernelAbstractions.jl.

For MadNLP, I need to to add an extension but we should have all LAPACK routines already available in oneAPI.jl.

[michel2323]

Let's hope they don't segfault.

[amontoison]

@michel2323 MadNLP/MadNLP.jl#476
Intel only supports CSR and COO. I will need to use a few linear algebra tricks to have the support of CSC.