mma_device.f90

 
 
submodule(mma) mma_device
 
  use device_math, only: device_copy, device_cmult, device_cadd, device_cfill, &
       device_add2, device_add3s2, device_invcol2, device_col2, device_col3, &
       device_sub2, device_sub3, device_add2s2, device_cadd2, device_pwmax2, &
       device_glsum, device_cmult2
  use device_mma_math, only: device_maxval, device_norm, device_lcsc2, &
       device_maxval2, device_maxval3, device_mma_gensub3, &
       device_mma_gensub4, device_mma_max, device_max2, device_rex, &
       device_relambda, device_delx, device_add2inv2, device_gg, device_diagx, &
       device_bb, device_updatebb, device_aa, device_updateaa, device_dx, &
       device_dy, device_dxsi, device_deta, device_kkt_rex, &
       device_mma_gensub2, device_mattrans_v_mul, device_mma_dipsolvesub1, &
       device_mma_ljjxinv, device_hess, device_solve_linear_system, &
       device_prepare_hessian, device_prepare_aa_matrix
 
  use neko_config, only: neko_bcknd_device, neko_device_mpi
  use device, only: device_to_host
  use comm, only: neko_comm, pe_rank, mpi_real_precision
  use mpi_f08, only: mpi_in_place, mpi_max, mpi_min
  use profiler, only: profiler_start_region, profiler_end_region
  use scratch_registry, only: neko_scratch_registry
 
  implicit none
 
contains
 
  module subroutine mma_update_device(this, iter, x, df0dx, fval, dfdx)
    ! ----------------------------------------------------- !
    ! Update the design variable x by solving the convex    !
    ! approximation of the problem.                         !
    !                                                       !
    ! This subroutine is called in each iteration of the    !
    ! optimization loop                                     !
    ! ----------------------------------------------------- !
    class(mma_t), intent(inout) :: this
    integer, intent(in) :: iter
    type(c_ptr), intent(inout) :: x
    type(c_ptr), intent(in) :: df0dx, fval, dfdx
 
    if (.not. this%is_initialized) then
       call neko_error("The MMA object is not initialized.")
    end if
 
    call profiler_start_region("MMA gensub")
    ! generate a convex approximation of the problem
    call mma_gensub_device(this, iter, x, df0dx, fval, dfdx)
    call profiler_end_region("MMA gensub")
 
    !solve the approximation problem using interior point method
    call profiler_start_region("MMA subsolve")
    if (this%subsolver .eq. "dip") then
       call mma_subsolve_dip_device(this, x)
    else if (this%subsolver .eq. "dpip") then
       call mma_subsolve_dpip_device(this, x)
    else
       call neko_error("Unrecognized subsolver for MMA in mma_device.")
    end if
    call profiler_end_region("MMA subsolve")
 
    this%is_updated = .true.
  end subroutine mma_update_device
 
  module subroutine mma_kkt_device(this, x, df0dx, fval, dfdx)
    class(mma_t), intent(inout) :: this
    type(c_ptr), intent(in) :: x, df0dx, fval, dfdx
 
    if (this%subsolver .eq. "dip") then
       call mma_dip_kkt_device(this, x, df0dx, fval, dfdx)
    else
       call mma_dpip_kkt_device(this, x, df0dx, fval, dfdx)
    end if
  end subroutine mma_kkt_device
 
  ! point method (dip) subsolve of MMA algorithm.
  module subroutine mma_dip_kkt_device(this, x, df0dx, fval, dfdx)
    class(mma_t), intent(inout) :: this
    type(c_ptr), intent(in) :: x, df0dx, fval, dfdx
 
    type(vector_t), pointer :: relambda, remu
    integer :: ind(2)
 
    call neko_scratch_registry%request(relambda, ind(1), this%m, .false.)
    call neko_scratch_registry%request(remu, ind(2), this%m, .false.)
 
    ! relambda = fval - this%a%x * this%z - this%y%x + this%mu%x
    call device_add3s2(relambda%x_d, fval, this%a%x_d, 1.0_rp, -this%z, &
         this%m)
    call device_sub2(relambda%x_d, this%y%x_d, this%m)
    call device_add2(relambda%x_d, this%mu%x_d, this%m)
 
    ! Compute residual for mu (eta in the paper)
    call device_col3(remu%x_d, this%lambda%x_d, this%mu%x_d, this%m)
 
    this%residumax = maxval([device_maxval(relambda%x_d, this%m), &
         device_maxval(remu%x_d, this%m)])
    this%residunorm = sqrt(device_norm(relambda%x_d, this%m)+ &
         device_norm(remu%x_d, this%m))
 
    call neko_scratch_registry%relinquish(ind)
  end subroutine mma_dip_kkt_device
 
  ! point method (dpip) subsolve of MMA algorithm.
  module subroutine mma_dpip_kkt_device(this, x, df0dx, fval, dfdx)
    class(mma_t), intent(inout) :: this
    type(c_ptr), intent(in) :: x, df0dx, fval, dfdx
 
    real(kind=rp) :: rez, rezeta
    type(vector_t), pointer :: rey, relambda, remu, res
    type(vector_t), pointer :: rex, rexsi, reeta
    integer :: ierr, ind(7)
    real(kind=rp) :: re_sq_norm
 
    call neko_scratch_registry%request(rey, ind(1), this%m, .false.)
    call neko_scratch_registry%request(relambda, ind(2), this%m, .false.)
    call neko_scratch_registry%request(remu, ind(3), this%m, .false.)
    call neko_scratch_registry%request(res, ind(4), this%m, .false.)
 
    call neko_scratch_registry%request(rex, ind(5), this%n, .false.)
    call neko_scratch_registry%request(rexsi, ind(6), this%n, .false.)
    call neko_scratch_registry%request(reeta, ind(7), this%n, .false.)
 
    call device_kkt_rex(rex%x_d, df0dx, dfdx, this%xsi%x_d, &
         this%eta%x_d, this%lambda%x_d, this%n, this%m)
 
    call device_col3(rey%x_d, this%d%x_d, this%y%x_d, this%m)
    call device_add2(rey%x_d, this%c%x_d, this%m)
    call device_sub2(rey%x_d, this%lambda%x_d, this%m)
    call device_sub2(rey%x_d, this%mu%x_d, this%m)
 
    rez = this%a0 - this%zeta - device_lcsc2(this%lambda%x_d, this%a%x_d, &
         this%m)
 
    call device_add3s2(relambda%x_d, fval, this%a%x_d, 1.0_rp, -this%z, &
         this%m)
    call device_sub2(relambda%x_d, this%y%x_d, this%m)
    call device_add2(relambda%x_d, this%s%x_d, this%m)
 
    call device_sub3(rexsi%x_d, x, this%xmin%x_d, this%n)
    call device_col2(rexsi%x_d, this%xsi%x_d, this%n)
 
    call device_sub3(reeta%x_d, this%xmax%x_d, x, this%n)
    call device_col2(reeta%x_d, this%eta%x_d, this%n)
 
    call device_col3(remu%x_d, this%mu%x_d, this%y%x_d, this%m)
 
    rezeta = this%zeta * this%z
 
    call device_col3(res%x_d, this%lambda%x_d, this%s%x_d, this%m)
 
    this%residumax = maxval([ &
         device_maxval(rex%x_d, this%n), &
         device_maxval(rey%x_d, this%m), &
         abs(rez), &
         device_maxval(relambda%x_d, this%m), &
         device_maxval(rexsi%x_d, this%n), &
         device_maxval(reeta%x_d, this%n), &
         device_maxval(remu%x_d, this%m), &
         abs(rezeta), &
         device_maxval(res%x_d, this%m)])
 
    re_sq_norm = device_norm(rex%x_d, this%n) + &
         device_norm(rexsi%x_d, this%n) + &
         device_norm(reeta%x_d, this%n)
 
    call mpi_allreduce(mpi_in_place, this%residumax, 1, &
         mpi_real_precision, mpi_max, neko_comm, ierr)
 
    call mpi_allreduce(mpi_in_place, re_sq_norm, 1, &
         mpi_real_precision, mpi_sum, neko_comm, ierr)
 
    this%residunorm = sqrt(( &
         device_norm(rey%x_d, this%m) + &
         rez**2 + &
         device_norm(relambda%x_d, this%m) + &
         device_norm(remu%x_d, this%m) + &
         rezeta**2 + &
         device_norm(res%x_d, this%m) &
         ) + re_sq_norm)
 
    call neko_scratch_registry%relinquish(ind)
  end subroutine mma_dpip_kkt_device
 
  !============================================================================!
  ! private internal subroutines
 
  subroutine mma_gensub_device(this, iter, x, df0dx, fval, dfdx)
    ! ----------------------------------------------------- !
    ! Generate the approximation sub problem by computing   !
    ! the lower and upper asymtotes and the other necessary !
    ! parameters (alpha, beta, p0j, q0j, pij, qij, ...).    !
    ! ----------------------------------------------------- !
    class(mma_t), intent(inout) :: this
    type(c_ptr), intent(in) :: x
    type(c_ptr), intent(in) :: df0dx
    type(c_ptr), intent(in) :: fval
    type(c_ptr), intent(in) :: dfdx
 
    integer, intent(in) :: iter
    integer :: ierr
 
    type(vector_t), pointer :: x_diff
    integer :: ind
 
    call neko_scratch_registry%request(x_diff, ind, this%n, .false.)
 
    call device_sub3(x_diff%x_d, this%xmax%x_d, this%xmin%x_d, this%n)
 
    ! ------------------------------------------------------------------------ !
    ! Setup the current asymptotes
 
    if (iter .lt. 3) then
       call device_copy(this%low%x_d, x, this%n)
       call device_add2s2(this%low%x_d, x_diff%x_d, - this%asyinit, this%n)
       call device_copy(this%upp%x_d, x, this%n)
       call device_add2s2(this%upp%x_d, x_diff%x_d, this%asyinit, this%n)
    else
       call device_mma_gensub2(this%low%x_d, this%upp%x_d, x, &
            this%xold1%x_d, this%xold2%x_d, x_diff%x_d, &
            this%asydecr, this%asyincr, this%n)
    end if
 
    ! ------------------------------------------------------------------------ !
    ! Calculate p0j, q0j, pij, qij, alpha, and beta
 
    call device_mma_gensub3(x, df0dx, dfdx, this%low%x_d, &
         this%upp%x_d, this%xmin%x_d, this%xmax%x_d, this%alpha%x_d, &
         this%beta%x_d, this%p0j%x_d, this%q0j%x_d, this%pij%x_d, &
         this%qij%x_d, this%n, this%m)
 
    ! ------------------------------------------------------------------------ !
    ! Computing bi as defined in page 5
 
    call device_mma_gensub4(x, this%low%x_d, this%upp%x_d, this%pij%x_d, &
         this%qij%x_d, this%n, this%m, this%bi%x_d)
 
    if (neko_device_mpi) then
       call mpi_allreduce(mpi_in_place, this%bi%x_d, this%m, &
            mpi_real_precision, mpi_sum, neko_comm, ierr)
    else
       call device_memcpy(this%bi%x, this%bi%x_d, this%m, device_to_host, &
            sync = .true.)
       call mpi_allreduce(mpi_in_place, this%bi%x, this%m, &
            mpi_real_precision, mpi_sum, neko_comm, ierr)
       call device_memcpy(this%bi%x, this%bi%x_d, this%m, host_to_device, &
            sync = .true.)
    end if
    call device_sub2(this%bi%x_d, fval, this%m)
 
    call neko_scratch_registry%relinquish(ind)
  end subroutine mma_gensub_device
 
  subroutine mma_subsolve_dpip_device(this, designx_d)
    class(mma_t), intent(inout) :: this
    type(c_ptr), intent(in) :: designx_d
    integer :: iter, itto, ierr
    real(kind=rp) :: epsi, residual_max, residual_norm, z, zeta, rez, rezeta, &
         delz, dz, dzeta, steg, zold, zetaold, new_residual
    ! vectors with size m
    type(vector_t) , pointer :: y, lambda, s, mu, rey, relambda, remu, res, &
         dely, dellambda, dy, dlambda, ds, dmu, yold, lambdaold, sold, muold
 
    ! vectors with size n
    type(vector_t), pointer :: x, xsi, eta, rex, rexsi, reeta, &
         delx, diagx, dx, dxsi, deta, xold, xsiold, etaold
 
    type(vector_t), pointer :: bb
    type(matrix_t), pointer :: GG
    type(matrix_t), pointer :: AA
 
    integer :: info
    real(kind=rp) :: re_sq_norm
 
    integer :: ind(35)
 
    real(kind=rp) :: minimal_epsilon
 
    call neko_scratch_registry%request(y, ind(1), this%m, .false.)
    call neko_scratch_registry%request(lambda, ind(2), this%m, .false.)
    call neko_scratch_registry%request(s, ind(3), this%m, .false.)
    call neko_scratch_registry%request(mu, ind(4), this%m, .false.)
    call neko_scratch_registry%request(rey, ind(5), this%m, .false.)
    call neko_scratch_registry%request(relambda, ind(6), this%m, .false.)
    call neko_scratch_registry%request(remu, ind(7), this%m, .false.)
    call neko_scratch_registry%request(res, ind(8), this%m, .false.)
    call neko_scratch_registry%request(dely, ind(9), this%m, .false.)
    call neko_scratch_registry%request(dellambda, ind(10), this%m, .false.)
    call neko_scratch_registry%request(dy, ind(11), this%m, .false.)
    call neko_scratch_registry%request(dlambda, ind(12), this%m, .false.)
    call neko_scratch_registry%request(ds, ind(13), this%m, .false.)
    call neko_scratch_registry%request(dmu, ind(14), this%m, .false.)
    call neko_scratch_registry%request(yold, ind(15), this%m, .false.)
    call neko_scratch_registry%request(lambdaold, ind(16), this%m, .false.)
    call neko_scratch_registry%request(sold, ind(17), this%m, .false.)
    call neko_scratch_registry%request(muold, ind(18), this%m, .false.)
    call neko_scratch_registry%request(x, ind(19), this%n, .false.)
    call neko_scratch_registry%request(xsi, ind(20), this%n, .false.)
    call neko_scratch_registry%request(eta, ind(21), this%n, .false.)
    call neko_scratch_registry%request(rex, ind(22), this%n, .false.)
    call neko_scratch_registry%request(rexsi, ind(23), this%n, .false.)
    call neko_scratch_registry%request(reeta, ind(24), this%n, .false.)
    call neko_scratch_registry%request(delx, ind(25), this%n, .false.)
    call neko_scratch_registry%request(diagx, ind(26), this%n, .false.)
    call neko_scratch_registry%request(dx, ind(27), this%n, .false.)
    call neko_scratch_registry%request(dxsi, ind(28), this%n, .false.)
    call neko_scratch_registry%request(deta, ind(29), this%n, .false.)
    call neko_scratch_registry%request(xold, ind(30), this%n, .false.)
    call neko_scratch_registry%request(xsiold, ind(31), this%n, .false.)
    call neko_scratch_registry%request(etaold, ind(32), this%n, .false.)
    call neko_scratch_registry%request(bb, ind(33), this%m+1, .false.)
 
    call neko_scratch_registry%request(gg, ind(34), this%m, this%n, .false.)
    call neko_scratch_registry%request(aa, ind(35), this%m+1, this%m+1, .false.)
 
    ! ------------------------------------------------------------------------ !
    ! initial value for the parameters in the subsolve based on
    ! page 15 of "https://people.kth.se/~krille/mmagcmma.pdf"
 
    epsi = 1.0_rp !100
    call device_add3s2(x%x_d, this%alpha%x_d, this%beta%x_d, 0.5_rp, 0.5_rp, &
         this%n)
    call device_cfill(y%x_d, 1.0_rp, this%m)
    z = 1.0_rp
    zeta = 1.0_rp
    call device_cfill(lambda%x_d, 1.0_rp, this%m)
    call device_cfill(s%x_d, 1.0_rp, this%m)
    call device_mma_max(xsi%x_d, x%x_d, this%alpha%x_d, this%n)
    call device_mma_max(eta%x_d, this%beta%x_d, x%x_d, this%n)
    call device_max2(mu%x_d, 1.0_rp, this%c%x_d, 0.5_rp, this%m)
 
    ! ------------------------------------------------------------------------ !
    ! Computing the minimal epsilon and choose the most conservative one
 
    minimal_epsilon = max(0.9_rp * this%epsimin, 1.0e-12_rp)
    call mpi_allreduce(mpi_in_place, minimal_epsilon, 1, &
         mpi_real_precision, mpi_min, neko_comm, ierr)
 
    ! ------------------------------------------------------------------------ !
    ! The main loop of the dual-primal interior point method.
 
    do while (epsi .gt. minimal_epsilon)
 
       ! --------------------------------------------------------------------- !
       ! Calculating residuals based on
       ! "https://people.kth.se/~krille/mmagcmma.pdf" for the variables
       ! x, y, z, lambda residuals based on eq(5.9a)-(5.9d), respectively.
 
       associate(p0j => this%p0j, q0j => this%q0j, &
            pij => this%pij, qij => this%qij, &
            low => this%low, upp => this%upp, &
            alpha => this%alpha, beta => this%beta, &
            c => this%c, d => this%d, &
            a0 => this%a0, a => this%a)
 
         call device_rex(rex%x_d, x%x_d, low%x_d, upp%x_d, &
              pij%x_d, p0j%x_d, qij%x_d, q0j%x_d, &
              lambda%x_d, xsi%x_d, eta%x_d, this%n, this%m)
 
         call device_col3(rey%x_d, d%x_d, y%x_d, this%m)
         call device_add2(rey%x_d, c%x_d, this%m)
         call device_sub2(rey%x_d, lambda%x_d, this%m)
         call device_sub2(rey%x_d, mu%x_d, this%m)
         rez = a0 - zeta - device_lcsc2(lambda%x_d, a%x_d, this%m)
 
         call device_cfill(relambda%x_d, 0.0_rp, this%m)
         call device_relambda(relambda%x_d, x%x_d, this%upp%x_d, &
              low%x_d, pij%x_d, qij%x_d, this%n, this%m)
 
       end associate
 
       ! --------------------------------------------------------------------- !
       ! Computing the norm of the residuals
 
       ! Complete the computations of lambda residuals
       if (neko_device_mpi) then
          call mpi_allreduce(mpi_in_place, relambda%x_d, this%m, &
               mpi_real_precision, mpi_sum, neko_comm, ierr)
       else
          call device_memcpy(relambda%x, relambda%x_d, this%m, device_to_host, &
               sync = .true.)
          call mpi_allreduce(mpi_in_place, relambda%x, this%m, &
               mpi_real_precision, mpi_sum, neko_comm, ierr)
          call device_memcpy(relambda%x, relambda%x_d, this%m, host_to_device, &
               sync = .true.)
       end if
 
       call device_add2s2(relambda%x_d, this%a%x_d, -z, this%m)
       call device_sub2(relambda%x_d, y%x_d, this%m)
       call device_add2(relambda%x_d, s%x_d, this%m)
       call device_sub2(relambda%x_d, this%bi%x_d, this%m)
 
       call device_sub3(rexsi%x_d, x%x_d, this%alpha%x_d, this%n)
       call device_col2(rexsi%x_d, xsi%x_d, this%n)
       call device_cadd(rexsi%x_d, - epsi, this%n)
 
       call device_sub3(reeta%x_d, this%beta%x_d, x%x_d, this%n)
       call device_col2(reeta%x_d, eta%x_d, this%n)
       call device_cadd(reeta%x_d, - epsi, this%n)
 
       call device_col3(remu%x_d, mu%x_d, y%x_d, this%m)
       call device_cadd(remu%x_d, - epsi, this%m)
 
       rezeta = zeta * z - epsi
 
       call device_col3(res%x_d, lambda%x_d, s%x_d, this%m)
       call device_cadd(res%x_d, - epsi, this%m)
 
       ! Setup vectors of residuals and their norms
       residual_max = maxval([device_maxval(rex%x_d, this%n), &
            device_maxval(rey%x_d, this%m), abs(rez), &
            device_maxval(relambda%x_d, this%m), &
            device_maxval(rexsi%x_d, this%n), &
            device_maxval(reeta%x_d, this%n), &
            device_maxval(remu%x_d, this%m), abs(rezeta), &
            device_maxval(res%x_d, this%m)])
 
       re_sq_norm = device_norm(rex%x_d, this%n) + &
            device_norm(rexsi%x_d, this%n) + device_norm(reeta%x_d, this%n)
 
       call mpi_allreduce(mpi_in_place, residual_max, 1, &
            mpi_real_precision, mpi_max, neko_comm, ierr)
 
       call mpi_allreduce(mpi_in_place, re_sq_norm, &
            1, mpi_real_precision, mpi_sum, neko_comm, ierr)
 
       residual_norm = sqrt(device_norm(rey%x_d, this%m) + &
            rez**2 + &
            device_norm(relambda%x_d, this%m) + &
            device_norm(remu%x_d, this%m)+ &
            rezeta**2 + &
            device_norm(res%x_d, this%m) &
            + re_sq_norm)
 
       ! --------------------------------------------------------------------- !
       ! Internal loop
 
       do iter = 1, this%max_iter
 
          if (residual_max .lt. epsi) exit
 
          call device_delx(delx%x_d, x%x_d, this%low%x_d, this%upp%x_d, &
               this%pij%x_d, this%qij%x_d, this%p0j%x_d, this%q0j%x_d, &
               this%alpha%x_d, this%beta%x_d, lambda%x_d, epsi, this%n, &
               this%m)
 
          call device_col3(dely%x_d, this%d%x_d, y%x_d, this%m)
          call device_add2(dely%x_d, this%c%x_d, this%m)
          call device_sub2(dely%x_d, lambda%x_d, this%m)
          call device_add2inv2(dely%x_d, y%x_d, - epsi, this%m)
          delz = this%a0 - device_lcsc2(lambda%x_d, this%a%x_d, this%m) - epsi/z
 
          ! Accumulate sums for dellambda (the term gi(x))
          call device_cfill(dellambda%x_d, 0.0_rp, this%m)
          call device_relambda(dellambda%x_d, x%x_d, this%upp%x_d, &
               this%low%x_d, this%pij%x_d, this%qij%x_d, this%n, this%m)
 
          call device_memcpy(dellambda%x, dellambda%x_d, this%m, &
               device_to_host, sync = .true.)
          call mpi_allreduce(mpi_in_place, dellambda%x, this%m, &
               mpi_real_precision, mpi_sum, neko_comm, ierr)
          call device_memcpy(dellambda%x, dellambda%x_d, this%m, &
               host_to_device, sync = .true.)
 
          call device_add3s2(dellambda%x_d, dellambda%x_d, this%a%x_d, &
               1.0_rp, -z, this%m)
          call device_sub2(dellambda%x_d, y%x_d, this%m)
          call device_sub2(dellambda%x_d, this%bi%x_d, this%m)
          call device_add2inv2(dellambda%x_d, lambda%x_d, epsi, this%m)
 
          call device_gg(gg%x_d, x%x_d, this%low%x_d, this%upp%x_d, &
               this%pij%x_d, this%qij%x_d, this%n, this%m)
 
          call device_diagx(diagx%x_d, x%x_d, xsi%x_d, this%low%x_d, &
               this%upp%x_d, this%p0j%x_d, this%q0j%x_d, this%pij%x_d, &
               this%qij%x_d, this%alpha%x_d, this%beta%x_d, eta%x_d, &
               lambda%x_d, this%n, this%m)
 
          !Here we only consider the case m<n in the matlab code
          !assembling the right hand side matrix based on eq(5.20)
          ! bb = [dellambda + dely/(this%d%x + &
          !         (mu/y)) - matmul(GG,delx/diagx), delz ]
 
          !--------------------------------------------------------------------!
          ! for MPI computation of bb
 
          call device_bb(bb%x_d, gg%x_d, delx%x_d, diagx%x_d, this%n, &
               this%m)
 
          call device_memcpy(bb%x, bb%x_d, this%m + 1, device_to_host, &
               sync = .true.)
          call mpi_allreduce(mpi_in_place, bb%x, this%m + 1, &
               mpi_real_precision, mpi_sum, neko_comm, ierr)
          call device_memcpy(bb%x, bb%x_d, this%m + 1, &
               host_to_device, sync = .true.)
 
          call device_updatebb(bb%x_d, dellambda%x_d, dely%x_d, &
               this%d%x_d, mu%x_d, y%x_d, delz, this%m)
 
          ! assembling the coefficients matrix AA based on eq(5.20)
          ! AA(1:this%m,1:this%m) =  &
          ! matmul(matmul(GG,mma_diag(1/diagx)), transpose(GG))
          ! !update diag(AA)
          ! AA(1:this%m,1:this%m) = AA(1:this%m,1:this%m) + &
          !     mma_diag(s/lambda + 1.0/(this%d%x + (mu/y)))
 
          call device_cfill(aa%x_d, 0.0_rp, (this%m+1) * (this%m+1))
          call device_aa(aa%x_d, gg%x_d, diagx%x_d, this%n, this%m)
 
          call device_memcpy(aa%x, aa%x_d, (this%m+1) * (this%m+1), &
               device_to_host, sync = .true.)
          call mpi_allreduce(mpi_in_place, aa%x, &
               (this%m + 1)**2, mpi_real_precision, mpi_sum, neko_comm, ierr)
          call device_memcpy(aa%x, aa%x_d, (this%m+1) * (this%m+1), &
               host_to_device, sync = .true.)
 
          call device_prepare_aa_matrix(aa%x_d, s%x_d, lambda%x_d, &
               this%d%x_d, mu%x_d, y%x_d, this%a%x_d, zeta, z, this%m)
 
          ! Device solve for the linear system
          call device_solve_linear_system(aa%x_d, bb%x_d, this%m + 1, info)
          if (info .ne. 0) then
             call neko_error("Linear solver failed on the device in  " // &
                  "mma_subsolve_dpip")
          end if
 
          call device_copy(dlambda%x_d, bb%x_d, this%m)
 
 
          !We need to write the last element of bb to dz so this is necessary
          call device_memcpy(bb%x, bb%x_d, this%m+1, device_to_host, &
               sync = .true.)
          dz = bb%x(this%m + 1)
 
 
          ! based on eq(5.19)
          call device_dx(dx%x_d, delx%x_d, diagx%x_d, gg%x_d, &
               dlambda%x_d, this%n, this%m)
          call device_dy(dy%x_d, dely%x_d, dlambda%x_d, this%d%x_d, &
               mu%x_d, y%x_d, this%m)
          call device_dxsi(dxsi%x_d, xsi%x_d, dx%x_d, x%x_d, &
               this%alpha%x_d, epsi, this%n)
          call device_deta(deta%x_d, eta%x_d, dx%x_d, x%x_d, &
               this%beta%x_d, epsi, this%n)
 
          call device_col3(dmu%x_d, mu%x_d, dy%x_d, this%m)
          call device_cmult(dmu%x_d, -1.0_rp, this%m)
          call device_cadd(dmu%x_d, epsi, this%m)
          call device_invcol2(dmu%x_d, y%x_d, this%m)
          call device_sub2(dmu%x_d, mu%x_d, this%m)
          dzeta = -zeta + (epsi - zeta * dz) / z
          call device_col3(ds%x_d, dlambda%x_d, s%x_d, this%m)
          call device_cmult(ds%x_d, -1.0_rp, this%m)
          call device_cadd(ds%x_d, epsi, this%m)
          call device_invcol2(ds%x_d, lambda%x_d, this%m)
          call device_sub2(ds%x_d, s%x_d, this%m)
 
          steg = maxval([1.0_rp, &
               device_maxval2(dy%x_d, y%x_d, -1.01_rp, this%m), &
               -1.01_rp * dz / z, &
               device_maxval2(dlambda%x_d, lambda%x_d, -1.01_rp, this%m), &
               device_maxval2(dxsi%x_d, xsi%x_d, -1.01_rp, this%n), &
               device_maxval2(deta%x_d, eta%x_d, -1.01_rp, this%n), &
               device_maxval2(dmu%x_d, mu%x_d, -1.01_rp, this%m), &
               -1.01_rp * dzeta / zeta, &
               device_maxval2(ds%x_d, s%x_d, -1.01_rp, this%m), &
               device_maxval3(dx%x_d, x%x_d, this%alpha%x_d, -1.01_rp, this%n),&
               device_maxval3(dx%x_d, this%beta%x_d, x%x_d, 1.01_rp, this%n)])
 
          steg = 1.0_rp / steg
 
          call device_copy(xold%x_d, x%x_d, this%n)
          call device_copy(yold%x_d, y%x_d, this%m)
          zold = z
          call device_copy(lambdaold%x_d, lambda%x_d, this%m)
          call device_copy(xsiold%x_d, xsi%x_d, this%n)
          call device_copy(etaold%x_d, eta%x_d, this%n)
          call device_copy(muold%x_d, mu%x_d, this%m)
          zetaold = zeta
          call device_copy(sold%x_d, s%x_d, this%m)
 
          new_residual = 2.0_rp * residual_norm
 
          ! Share the new_residual and steg values
          call mpi_allreduce(mpi_in_place, steg, 1, &
               mpi_real_precision, mpi_min, neko_comm, ierr)
          call mpi_allreduce(mpi_in_place, new_residual, 1, &
               mpi_real_precision, mpi_min, neko_comm, ierr)
 
          ! The innermost loop to determine the suitable step length
          ! using the Backtracking Line Search approach
          itto = 0
          do while ((new_residual .gt. residual_norm) .and. (itto .lt. 50))
             itto = itto + 1
 
             ! update the variables
             call device_add3s2(x%x_d, xold%x_d, dx%x_d, 1.0_rp, steg, this%n)
             call device_add3s2(y%x_d, yold%x_d, dy%x_d, 1.0_rp, steg, this%m)
             z = zold + steg*dz
             call device_add3s2(lambda%x_d, lambdaold%x_d, &
                  dlambda%x_d, 1.0_rp, steg, this%m)
             call device_add3s2(xsi%x_d, xsiold%x_d, dxsi%x_d, &
                  1.0_rp, steg, this%n)
             call device_add3s2(eta%x_d, etaold%x_d, deta%x_d, &
                  1.0_rp, steg, this%n)
             call device_add3s2(mu%x_d, muold%x_d, dmu%x_d, &
                  1.0_rp, steg, this%m)
             zeta = zetaold + steg*dzeta
             call device_add3s2(s%x_d, sold%x_d, ds%x_d, 1.0_rp, &
                  steg, this%m)
 
             ! Recompute the new_residual to see if this stepsize improves
             ! the residue
             call device_rex(rex%x_d, x%x_d, this%low%x_d, &
                  this%upp%x_d, this%pij%x_d, this%p0j%x_d, &
                  this%qij%x_d, this%q0j%x_d, lambda%x_d, xsi%x_d, &
                  eta%x_d, this%n, this%m)
 
             call device_col3(rey%x_d, this%d%x_d, y%x_d, this%m)
             call device_add2(rey%x_d, this%c%x_d, this%m)
             call device_sub2(rey%x_d, lambda%x_d, this%m)
             call device_sub2(rey%x_d, mu%x_d, this%m)
 
             rez = this%a0 - zeta - device_lcsc2(lambda%x_d, this%a%x_d, this%m)
 
             ! Accumulate sums for relambda (the term gi(x))
             call device_cfill(relambda%x_d, 0.0_rp, this%m)
             call device_relambda(relambda%x_d, x%x_d, this%upp%x_d, &
                  this%low%x_d, this%pij%x_d, this%qij%x_d, &
                  this%n, this%m)
 
             call device_memcpy(relambda%x, relambda%x_d, this%m, &
                  device_to_host, sync = .true.)
             call mpi_allreduce(mpi_in_place, relambda%x, this%m, &
                  mpi_real_precision, mpi_sum, neko_comm, ierr)
             call device_memcpy(relambda%x, relambda%x_d, &
                  this%m, host_to_device, sync = .true.)
 
             call device_add3s2(relambda%x_d, relambda%x_d, &
                  this%a%x_d, 1.0_rp, -z, this%m)
             call device_sub2(relambda%x_d, y%x_d, this%m)
             call device_add2(relambda%x_d, s%x_d, this%m)
             call device_sub2(relambda%x_d, this%bi%x_d, this%m)
 
             call device_sub3(rexsi%x_d, x%x_d, this%alpha%x_d, this%n)
             call device_col2(rexsi%x_d, xsi%x_d, this%n)
             call device_cadd(rexsi%x_d, - epsi, this%n)
 
             call device_sub3(reeta%x_d, this%beta%x_d, x%x_d, this%n)
             call device_col2(reeta%x_d, eta%x_d, this%n)
             call device_cadd(reeta%x_d, - epsi, this%n)
 
             call device_col3(remu%x_d, mu%x_d, y%x_d, this%m)
             call device_cadd(remu%x_d, - epsi, this%m)
 
             rezeta = zeta*z - epsi
 
             call device_col3(res%x_d, lambda%x_d, s%x_d, this%m)
             call device_cadd(res%x_d, - epsi, this%m)
 
             ! Compute squared norms for the residuals
             re_sq_norm = device_norm(rex%x_d, this%n) + &
                  device_norm(rexsi%x_d, this%n) + &
                  device_norm(reeta%x_d, this%n)
             call mpi_allreduce(mpi_in_place, re_sq_norm, 1, &
                  mpi_real_precision, mpi_sum, neko_comm, ierr)
 
             new_residual = sqrt(device_norm(rey%x_d, this%m) + &
                  rez**2 + &
                  device_norm(relambda%x_d, this%m) + &
                  device_norm(remu%x_d, this%m) + &
                  rezeta**2 + &
                  device_norm(res%x_d, this%m) + &
                  re_sq_norm)
 
             steg = steg / 2.0_rp
 
          end do
          steg = 2.0_rp * steg ! Correction for the final division by 2
 
          ! Update the maximum and norm of the residuals
          residual_norm = new_residual
          residual_max = maxval([ &
               device_maxval(rex%x_d, this%n), &
               device_maxval(rey%x_d, this%m), &
               abs(rez), &
               device_maxval(relambda%x_d, this%m), &
               device_maxval(rexsi%x_d, this%n), &
               device_maxval(reeta%x_d, this%n), &
               device_maxval(remu%x_d, this%m), &
               abs(rezeta), &
               device_maxval(res%x_d, this%m)])
 
          call mpi_allreduce(mpi_in_place, residual_max, 1, &
               mpi_real_precision, mpi_max, neko_comm, ierr)
 
       end do
 
       epsi = 0.1_rp * epsi
    end do
 
    ! Save the new designx
    call device_copy(this%xold2%x_d, this%xold1%x_d, this%n)
    call device_copy(this%xold1%x_d, designx_d, this%n)
    call device_copy(designx_d, x%x_d, this%n)
 
    ! update the parameters of the MMA object nesessary to compute KKT residual
    call device_copy(this%y%x_d, y%x_d, this%m)
    this%z = z
    call device_copy(this%lambda%x_d, lambda%x_d, this%m)
    this%zeta = zeta
    call device_copy(this%xsi%x_d, xsi%x_d, this%n)
    call device_copy(this%eta%x_d, eta%x_d, this%n)
    call device_copy(this%mu%x_d, mu%x_d, this%m)
    call device_copy(this%s%x_d, s%x_d, this%m)
 
    !free all the initiated variables in this subroutine
    call neko_scratch_registry%relinquish(ind)
  end subroutine mma_subsolve_dpip_device
 
  subroutine mma_subsolve_dip_device(this, designx_d)
    class(mma_t), intent(inout) :: this
    type(c_ptr), intent(in) :: designx_d
    integer :: iter, ierr
    real(kind=rp) :: epsi, residumax, z, steg
    ! vectors with size m
    type(vector_t), pointer :: y, lambda, mu, relambda, remu, dlambda, dmu, &
         gradlambda, zerom, dd, dummy_m
    ! vectors with size n
    type(vector_t), pointer :: x, pjlambda, qjlambda
 
    ! inverse of a diag matrix:
    type(vector_t), pointer :: Ljjxinv ! [∇_x^2 Ljj]−1
    type(matrix_t), pointer :: hijx ! ∇_x hij
    type(matrix_t), pointer :: Hess
 
    integer :: info, ind(17)
 
    real(kind=rp) :: minimal_epsilon
 
    call neko_scratch_registry%request(y, ind(1), this%m, .false.)
    call neko_scratch_registry%request(lambda, ind(2), this%m, .false.)
    call neko_scratch_registry%request(mu, ind(3), this%m, .false.)
    call neko_scratch_registry%request(relambda, ind(4), this%m, .false.)
    call neko_scratch_registry%request(remu, ind(5), this%m, .false.)
    call neko_scratch_registry%request(dlambda, ind(6), this%m, .false.)
    call neko_scratch_registry%request(dmu, ind(7), this%m, .false.)
    call neko_scratch_registry%request(gradlambda, ind(8), this%m, .false.)
    call neko_scratch_registry%request(zerom, ind(9), this%m, .false.)
    call neko_scratch_registry%request(dd, ind(10), this%m, .false.)
    call neko_scratch_registry%request(dummy_m, ind(11), this%m, .false.)
 
    call neko_scratch_registry%request(x, ind(12), this%n, .false.)
    call neko_scratch_registry%request(pjlambda,ind(13), this%n, .false.)
    call neko_scratch_registry%request(qjlambda, ind(14), this%n, .false.)
 
    call neko_scratch_registry%request(ljjxinv, ind(15), this%n, .false.)
 
    call neko_scratch_registry%request(hijx, ind(16), this%m, this%n, .false.)
    call neko_scratch_registry%request(hess, ind(17), this%m, this%m, .false.)
 
    ! ------------------------------------------------------------------------ !
    ! initial value for the parameters in the subsolve based on
    ! page 15 of "https://people.kth.se/~krille/mmagcmma.pdf"
 
    epsi = 1.0_rp !100
    call device_cfill(y%x_d, 1.0_rp, this%m)
    ! initialize lambda with an array of ones (change to this%c%x/2 if needed!)
    call device_cfill(lambda%x_d, 1.0_rp, this%m)
    call device_cmult2(dummy_m%x_d, this%c%x_d, 0.5_rp, this%m)
    call device_pwmax2(lambda%x_d, dummy_m%x_d, this%m)
 
    call device_cfill(mu%x_d, 1.0_rp, this%m)
    z = 0.0_rp
 
    ! dd is defined as this%d + 1.0e-8_rp, to avoid devision by 0 in computing y
    call device_cadd2(dd%x_d, this%d%x_d, 1.0e-8_rp, this%m)
 
    ! ------------------------------------------------------------------------ !
    ! Computing the minimal epsilon and choose the most conservative one
 
    minimal_epsilon = max(0.9_rp * this%epsimin, 1.0e-12_rp)
    call mpi_allreduce(mpi_in_place, minimal_epsilon, 1, &
         mpi_real_precision, mpi_min, neko_comm, ierr)
 
    ! ------------------------------------------------------------------------ !
    ! The main loop of the dual-primal interior point method.
 
    outer: do while (epsi .gt. minimal_epsilon)
       ! calculating residuals based on
       ! "https://people.kth.se/~krille/mmagcmma.pdf" for the variables
       ! x, y, z, lambda residuals based on eq(5.9a)-(5.9d), respectively.
       associate(p0j => this%p0j, q0j => this%q0j, &
            pij => this%pij, qij => this%qij, &
            low => this%low, upp => this%upp, &
            alpha => this%alpha, beta => this%beta, &
            c => this%c, a0 => this%a0, a => this%a)
 
         ! minimize(L_x, L_y, L_z) and compute x(λ), y(λ), z(λ) for
         ! the initial value of λ
 
         ! Comput the value of y that minimizes L_y for the current λ
         ! minimize (sum_{i=1}^{m} [ (c_i - λ_i) * y_i + 0.5 * d_i * y_i^2 ])
         ! dL_y/dy =0   => y= (λ_i - c_i)/d_i, ensure y>=0
         call device_sub3(y%x_d, lambda%x_d, c%x_d, this%m)
         ! division by dd to avoid devision by 0 (in case this%d%x_d = 0)
         call device_invcol2(y%x_d, dd%x_d, this%m)
         call device_pwmax2(y%x_d, zerom%x_d, this%m)
 
         ! Comput the value of z that minimizes L_z for the current λ
         ! minimize ((a_0 - sum_{i=1}^{m} λ_i * a_i) * z)
         ! if (a_0-dot_product(lambda, a)>=0) z=0 else z= 1.0
         ! ensure z>=0
         call device_col3(dummy_m%x_d, lambda%x_d, a%x_d, this%m)
         z = device_glsum(dummy_m%x_d, this%m)
         z = merge(0.0_rp, 1.0_rp, a0 - z >= 0.0)
 
         ! Comput the value of x that minimizes L_x for the current λ
         ! minimize( sum_{j=1}^{n} [ (p_{0j} + sum_{i=1}^{m} λ_i *
         ! p_{ij}) / (u_j - x_j) + (q_{0j} + sum_{i=1}^{m} λ_i * q_{ij}) /
         ! (x_j - l_j) ] - sum_{i=1}^{m} λ_i * b_i)
         call device_mattrans_v_mul(pjlambda%x_d, pij%x_d, lambda%x_d, this%m, this%n)
         call device_mattrans_v_mul(qjlambda%x_d, qij%x_d, lambda%x_d, this%m, this%n)
         call device_add2(pjlambda%x_d, p0j%x_d, this%n)
         call device_add2(qjlambda%x_d, q0j%x_d, this%n)
 
         call device_mma_dipsolvesub1(x%x_d, pjlambda%x_d, qjlambda%x_d, &
              low%x_d, upp%x_d, alpha%x_d, beta%x_d, this%n)
 
         call device_cfill(relambda%x_d, 0.0_rp, this%m)
         call device_relambda(relambda%x_d, x%x_d, this%upp%x_d, &
              low%x_d, pij%x_d, qij%x_d, this%n, this%m)
 
         ! Global comminucation for relambda values
 
         call device_memcpy(relambda%x, relambda%x_d, this%m, device_to_host, &
              sync = .true.)
         call mpi_allreduce(mpi_in_place, relambda%x, this%m, &
              mpi_real_precision, mpi_sum, neko_comm, ierr)
         call device_memcpy(relambda%x, relambda%x_d, this%m, &
              host_to_device, sync = .true.)
 
         call device_add2s2(relambda%x_d, this%a%x_d, -z, this%m)
         call device_sub2(relambda%x_d, y%x_d, this%m)
         call device_add2(relambda%x_d, mu%x_d, this%m)
         call device_sub2(relambda%x_d, this%bi%x_d, this%m)
 
         call device_col3(remu%x_d, mu%x_d, lambda%x_d, this%m)
         call device_cadd(remu%x_d, -epsi, this%m)
 
         residumax = maxval([device_maxval(relambda%x_d, this%m), &
              device_maxval(remu%x_d, this%m)])
 
         ! ------------------------------------------------------------------- !
         ! Internal loop
         do iter = 1, this%max_iter
            !Check the condition
            if (residumax .lt. epsi) exit
 
            ! Compute dL(x, y, z, λ)/dλ for the updated x(λ), y(λ), z(λ)
            ! based on the implementation in the following paper by Niels
            ! https://doi.org/10.1007/s00158-012-0869-2
            ! (https://github.com/topopt/TopOpt_in_PETSc/blob/master/MMA.cc)
            ! The formula for gradlambda and relambda are basically the same:
            ! thus, we utilise gradlambda = relambda - mu for efficiency
            call device_copy(gradlambda%x_d, relambda%x_d, this%m)
            call device_sub2(gradlambda%x_d, mu%x_d, this%m)
 
            ! Update gradlambda as the right hand side for Newton's method(eq10)
            call device_cfill(dummy_m%x_d, epsi, this%m)
            call device_invcol2(dummy_m%x_d, lambda%x_d, this%m)
            call device_add2(gradlambda%x_d, dummy_m%x_d, this%m)
            call device_cmult(gradlambda%x_d, -1.0_rp, this%m)
 
            ! Computing the Hessian as in equation (13) in
            !! https://doi.org/10.1007/s00158-012-0869-2
 
            !--------------contributions of x terms to Hess--------------------!
            call device_mma_ljjxinv(ljjxinv%x_d, pjlambda%x_d, qjlambda%x_d, &
                 x%x_d, low%x_d, upp%x_d, alpha%x_d, beta%x_d, this%n)
 
            call device_gg(hijx%x_d, x%x_d, this%low%x_d, this%upp%x_d, &
                 this%pij%x_d, this%qij%x_d, this%n, this%m)
 
            call device_cfill(hess%x_d, 0.0_rp, (this%m) * (this%m) )
            call device_hess(hess%x_d, hijx%x_d, ljjxinv%x_d, this%n, this%m)
 
            ! download Hess to CPU, mpi reduce, upload to the device
            call device_memcpy(hess%x, hess%x_d, this%m*this%m, device_to_host, &
                 sync = .true.)
            call mpi_allreduce(mpi_in_place, hess%x, &
                 this%m*this%m, mpi_real_precision, mpi_sum, neko_comm, ierr)
            ! No need to upload to device since we solve LSE on CPU
            ! But now we solve LSE on GPU, so upload it:
            call device_memcpy(hess%x, hess%x_d, this%m*this%m, &
                 host_to_device, sync = .true.)
 
            !---------------contributions of z terms to Hess-------------------!
            ! There is no contibution to the Hess from z terms as z terms are
            ! linear w.r.t λ
 
 
            !---------------contributions of y terms to Hess-------------------!
            ! Only for inactive constraint, we consider contributions to Hess.
            ! Note that if d(i) = 0, the y terms (just like z terms) will not
            ! contribute to the Hessian matrix.
            ! Note that since we use DGESV to solve LSE on CPU, we dont need
            ! cuda kernel for this part
            ! Also, improve the robustness by stablizing the Hess using
            ! Levenberg-Marquardt algorithm (heuristically)
            call device_prepare_hessian(hess%x_d, y%x_d, this%d%x_d, &
                 mu%x_d, lambda%x_d, this%m)
 
            ! Device solve for the linear system
            call device_solve_linear_system(hess%x_d, gradlambda%x_d, &
                 this%m, info)
            if (info .ne. 0) then
               call neko_error("Linear solver failed on the device in  " // &
                    "mma_subsolve_dip")
            end if
 
            call device_copy(dlambda%x_d, gradlambda%x_d, this%m)
 
            ! based on eq(11) for delta eta
            call device_copy(dummy_m%x_d, dlambda%x_d, this%m)
            call device_col2(dummy_m%x_d, mu%x_d, this%m)
            call device_invcol2(dummy_m%x_d, lambda%x_d, this%m)
 
            call device_cfill(dmu%x_d, epsi, this%m)
            call device_invcol2(dmu%x_d, lambda%x_d, this%m)
            call device_add2s2(dmu%x_d, dummy_m%x_d, -1.0_rp, this%m)
            call device_sub2(dmu%x_d, mu%x_d, this%m)
 
            steg = maxval([1.005_rp, device_maxval2(dlambda%x_d, lambda%x_d, &
                 -1.01_rp, this%m), device_maxval2(dmu%x_d, mu%x_d, -1.01_rp, &
                 this%m)])
            steg = 1.0_rp / steg
 
            call device_add2s2(lambda%x_d, dlambda%x_d, steg, this%m)
            call device_add2s2(mu%x_d, dmu%x_d, steg, this%m)
 
            ! minimize(L_x, L_y, L_z) and compute x(λ), y(λ), z(λ) for
            ! the updated values of λ
 
            ! Comput the value of y that minimizes L_y for the current λ
            ! minimize (sum_{i=1}^{m} [ (c_i - λ_i) * y_i + 0.5 * d_i * y_i^2 ])
            ! dL_y/dy =0   => y= (λ_i - c_i)/d_i, ensure y>=0
 
            call device_sub3(y%x_d, lambda%x_d, c%x_d, this%m)
            ! division by dd to avoid devision by 0 (in case this%d%x_d = 0)
            call device_invcol2(y%x_d, dd%x_d, this%m)
            call device_pwmax2(y%x_d, zerom%x_d, this%m)
 
            ! Comput the value of z that minimizes L_z for the current λ
            ! minimize ((a_0 - sum_{i=1}^{m} λ_i * a_i) * z)
            ! if (a_0-dot_product(lambda, a)>=0) z=0 else z= 1.0
            ! ensure z>=0
            call device_col3(dummy_m%x_d, lambda%x_d, a%x_d, this%m)
            z = device_glsum(dummy_m%x_d, this%m)
            z = merge(0.0_rp, 1.0_rp, a0 - z >= 0.0)
 
            ! Comput the value of x that minimizes L_x for the current λ
            ! minimize( sum_{j=1}^{n} [ (p_{0j} + sum_{i=1}^{m} λ_i *
            ! p_{ij}) / (u_j - x_j) + (q_{0j} + sum_{i=1}^{m} λ_i * q_{ij}) /
            ! (x_j - l_j) ] - sum_{i=1}^{m} λ_i * b_i)
            call device_mattrans_v_mul(pjlambda%x_d, pij%x_d, lambda%x_d, this%m, this%n)
            call device_mattrans_v_mul(qjlambda%x_d, qij%x_d, lambda%x_d, this%m, this%n)
            call device_add2(pjlambda%x_d, p0j%x_d, this%n)
            call device_add2(qjlambda%x_d, q0j%x_d, this%n)
 
            call device_mma_dipsolvesub1(x%x_d, pjlambda%x_d, qjlambda%x_d, &
                 low%x_d, upp%x_d, alpha%x_d, beta%x_d, this%n)
 
            ! Compute the residual for the lambda and mu using eq(9) and eq(15)
 
            call device_cfill(relambda%x_d, 0.0_rp, this%m)
            call device_relambda(relambda%x_d, x%x_d, this%upp%x_d, &
                 low%x_d, pij%x_d, qij%x_d, this%n, this%m)
 
            ! Global comminucation for relambda values
 
            call device_memcpy(relambda%x, relambda%x_d, this%m, device_to_host, &
                 sync = .true.)
            call mpi_allreduce(mpi_in_place, relambda%x, this%m, &
                 mpi_real_precision, mpi_sum, neko_comm, ierr)
            call device_memcpy(relambda%x, relambda%x_d, this%m, &
                 host_to_device, sync = .true.)
 
            call device_add2s2(relambda%x_d, this%a%x_d, -z, this%m)
            call device_sub2(relambda%x_d, y%x_d, this%m)
            call device_add2(relambda%x_d, mu%x_d, this%m)
            call device_sub2(relambda%x_d, this%bi%x_d, this%m)
 
            call device_col3(remu%x_d, mu%x_d, lambda%x_d, this%m)
            call device_cadd(remu%x_d, -epsi, this%m)
 
            residumax = maxval([device_maxval(relambda%x_d, this%m), &
                 device_maxval(remu%x_d, this%m)])
         end do
       end associate
       epsi = 0.1_rp * epsi
    end do outer
 
    ! Save the new designx
    call device_copy(this%xold2%x_d, this%xold1%x_d, this%n)
    call device_copy(this%xold1%x_d, designx_d, this%n)
    call device_copy(designx_d, x%x_d, this%n)
 
    ! update the parameters of the MMA object nesessary to compute KKT residual
    call device_copy(this%y%x_d, y%x_d, this%m)
    this%z = z
    call device_copy(this%lambda%x_d, lambda%x_d, this%m)
    call device_copy(this%mu%x_d, mu%x_d, this%m)
 
    call neko_scratch_registry%relinquish(ind)
  end subroutine mma_subsolve_dip_device
 
end submodule mma_device