71 use num_types,
only: rp, dp, sp
72 use json_module,
only: json_file
73 use json_utils,
only: json_get_or_default
74 use vector,
only: vector_t
75 use matrix,
only: matrix_t
76 use comm,
only: pe_rank, neko_comm, pe_size, mpi_real_precision
77 use utils,
only: neko_error, filename_suffix
78 use neko_config,
only: neko_bcknd_device, neko_bcknd_cuda, neko_bcknd_hip, &
80 use device,
only: device_memcpy, host_to_device, device_to_host
81 use,
intrinsic :: iso_c_binding, only: c_ptr
82 use logger,
only: neko_log
83 use mpi_f08,
only: mpi_sum, mpi_allreduce, mpi_integer
84 use scratch_registry,
only: scratch_registry_t
85 use profiler,
only: profiler_start_region, profiler_end_region
93 integer :: n, m, n_global, max_iter
94 real(kind=rp) :: a0, asyinit, asyincr, asydecr, epsimin, &
96 real(kind=rp) :: move_limit = 0.2_rp
97 type(vector_t) :: xold1, xold2, low, upp, alpha, beta, a, c, d, xmax, xmin
98 logical :: is_initialized = .false.
99 logical :: is_updated = .false.
100 type(scratch_registry_t) :: scratch
101 character(len=:),
allocatable :: subsolver, bcknd
104 type(vector_t) :: p0j, q0j
105 type(matrix_t) :: pij, qij
109 real(kind=rp) :: z, zeta
110 type(vector_t) :: y, lambda, s, mu
111 type(vector_t) :: xsi, eta
114 generic,
public :: init => init_from_json, init_from_components
115 procedure,
public, pass(this) :: init_from_json => mma_init_from_json
116 procedure,
public, pass(this) :: init_from_components => &
117 mma_init_from_components
118 procedure,
public, pass(this) :: free => mma_free
119 procedure,
public, pass(this) :: get_n => mma_get_n
120 procedure,
public, pass(this) :: get_m => mma_get_m
121 procedure,
public, pass(this) :: get_residumax => mma_get_residumax
122 procedure,
public, pass(this) :: get_residunorm => mma_get_residunorm
123 procedure,
public, pass(this) :: get_max_iter => mma_get_max_iter
124 procedure,
public, pass(this) :: get_backend_and_subsolver => &
125 mma_get_backend_and_subsolver
127 generic,
public :: update => update_vector, update_cpu, update_device
128 procedure, pass(this) :: update_vector => mma_update_vector
129 procedure, pass(this) :: update_cpu => mma_update_cpu
130 procedure, pass(this) :: update_device => mma_update_device
132 generic,
public :: kkt => kkt_vector, kkt_cpu, kkt_device
133 procedure, pass(this) :: kkt_vector => mma_kkt_vector
134 procedure, pass(this) :: kkt_cpu => mma_kkt_cpu
135 procedure, pass(this) :: kkt_device => mma_kkt_device
137 procedure, pass(this) :: save_checkpoint => mma_save_checkpoint
138 procedure, pass(this) :: load_checkpoint => mma_load_checkpoint
141 procedure, pass(this) :: copy_from => mma_copy_from
148 real(kind=rp),
parameter :: a0_default = 1.0_rp
149 real(kind=rp),
parameter :: a_default = 0.0_rp
150 real(kind=rp),
parameter :: c_default = 1000.0_rp
151 real(kind=rp),
parameter :: d_default = 0.0_rp
152 real(kind=rp),
parameter :: xmin_default = 0.0_rp
153 real(kind=rp),
parameter :: xmax_default = 1.0_rp
155 real(kind=rp),
parameter :: asyinit_default = 0.2_rp
156 real(kind=rp),
parameter :: asyincr_default = 1.05_rp
157 real(kind=rp),
parameter :: asydecr_default = 0.65_rp
158 real(kind=rp),
parameter :: move_limit_default = 0.2_rp
160 integer,
parameter :: max_iter_default = 100
161 character(len=*),
parameter :: subsolver_default =
"dip"
162 real(kind=rp),
parameter :: scale_default = 1.0_rp
163 logical,
parameter :: auto_scale_default = .false.
170 module subroutine mma_update_cpu(this, iter, x, df0dx, fval, dfdx)
171 class(mma_t),
intent(inout) :: this
172 integer,
intent(in) :: iter
173 real(kind=rp),
dimension(this%n),
intent(inout) :: x
174 real(kind=rp),
dimension(this%n),
intent(in) :: df0dx
175 real(kind=rp),
dimension(this%m),
intent(in) :: fval
176 real(kind=rp),
dimension(this%m, this%n),
intent(in) :: dfdx
177 end subroutine mma_update_cpu
180 module subroutine mma_kkt_cpu(this, x, df0dx, fval, dfdx)
181 class(mma_t),
intent(inout) :: this
182 real(kind=rp),
dimension(this%n),
intent(in) :: x
183 real(kind=rp),
dimension(this%n),
intent(in) :: df0dx
184 real(kind=rp),
dimension(this%m),
intent(in) :: fval
185 real(kind=rp),
dimension(this%m, this%n),
intent(in) :: dfdx
186 end subroutine mma_kkt_cpu
192 module subroutine mma_update_device(this, iter, x, df0dx, fval, dfdx)
193 class(mma_t),
intent(inout) :: this
194 integer,
intent(in) :: iter
195 type(c_ptr),
intent(inout) :: x
196 type(c_ptr),
intent(in) :: df0dx, fval, dfdx
197 end subroutine mma_update_device
200 module subroutine mma_kkt_device(this, x, df0dx, fval, dfdx)
201 class(mma_t),
intent(inout) :: this
202 type(c_ptr),
intent(in) :: x, df0dx, fval, dfdx
203 end subroutine mma_kkt_device
211 module subroutine mma_save_checkpoint_hdf5(object, filename, overwrite)
212 class(mma_t),
intent(inout) :: object
213 character(len=*),
intent(in) :: filename
214 logical,
intent(in),
optional :: overwrite
215 end subroutine mma_save_checkpoint_hdf5
217 module subroutine mma_load_checkpoint_hdf5(object, filename)
218 class(mma_t),
intent(inout) :: object
219 character(len=*),
intent(in) :: filename
220 end subroutine mma_load_checkpoint_hdf5
229 subroutine mma_init_from_json(this, x, n, m, json, scale, auto_scale)
241 class(mma_t),
intent(inout) :: this
242 integer,
intent(in) :: n, m
243 type(vector_t),
intent(in) :: x
245 type(json_file),
intent(inout) :: json
248 real(kind=rp),
intent(out) :: scale
249 logical,
intent(out) :: auto_scale
257 real(kind=rp),
dimension(n) :: xmax, xmin
258 real(kind=rp),
dimension(m) :: a, c, d
259 character(len=:),
allocatable :: subsolver, bcknd, bcknd_default
262 real(kind=rp) :: a0 , xmax_const, xmin_const, a_const, c_const, d_const
263 real(kind=rp) :: move_limit
265 integer :: max_iter, n_global, ierr
266 real(kind=rp) :: epsimin, asyinit, asyincr, asydecr
268 call mpi_allreduce(n, n_global, 1, mpi_integer, &
269 mpi_sum, neko_comm, ierr)
272 if (neko_bcknd_device .eq. 1)
then
273 bcknd_default =
"device"
275 bcknd_default =
"cpu"
281 call json_get_or_default(json,
'mma.epsimin', epsimin, &
282 1.0e-9_rp * sqrt(real(m + n_global, rp)))
283 call json_get_or_default(json,
'mma.max_iter', max_iter, max_iter_default)
286 call json_get_or_default(json,
'mma.asyinit', asyinit, asyinit_default)
287 call json_get_or_default(json,
'mma.asyincr', asyincr, asyincr_default)
288 call json_get_or_default(json,
'mma.asydecr', asydecr, asydecr_default)
290 call json_get_or_default(json,
'mma.backend', bcknd, bcknd_default)
291 call json_get_or_default(json,
'mma.subsolver', subsolver, subsolver_default)
293 call json_get_or_default(json,
'mma.xmin', xmin_const, xmin_default)
294 call json_get_or_default(json,
'mma.xmax', xmax_const, xmax_default)
295 call json_get_or_default(json,
'mma.a0', a0, a0_default)
296 call json_get_or_default(json,
'mma.a', a_const, a_default)
297 call json_get_or_default(json,
'mma.c', c_const, c_default)
298 call json_get_or_default(json,
'mma.d', d_const, d_default)
299 call json_get_or_default(json,
'mma.move_limit', move_limit, move_limit_default)
301 call json_get_or_default(json,
'mma.scale', scale, scale_default)
302 call json_get_or_default(json,
'mma.auto_scale', auto_scale, auto_scale_default)
314 call this%init(x, n, m, a0, a, c, d, xmin, xmax, &
315 max_iter, epsimin, asyinit, asyincr, asydecr, bcknd, subsolver, &
318 end subroutine mma_init_from_json
321 subroutine mma_free(this)
322 class(mma_t),
intent(inout) :: this
324 call this%xold1%free()
325 call this%xold2%free()
326 call this%alpha%free()
327 call this%beta%free()
333 call this%xmax%free()
334 call this%xmin%free()
339 call this%lambda%free()
348 call this%scratch%free()
350 this%is_initialized = .false.
351 this%is_updated = .false.
352 end subroutine mma_free
355 subroutine mma_init_from_components(this, x, n, m, a0, a, c, d, xmin, xmax, &
356 max_iter, epsimin, asyinit, asyincr, asydecr, bcknd, subsolver, &
369 class(mma_t),
intent(inout) :: this
370 integer,
intent(in) :: n, m
371 type(vector_t),
intent(in) :: x
379 real(kind=rp),
intent(in),
dimension(n) :: xmax, xmin
380 real(kind=rp),
intent(in),
dimension(m) :: a, c, d
381 real(kind=rp),
intent(in) :: a0
382 integer,
intent(in),
optional :: max_iter
383 real(kind=rp),
intent(in),
optional :: epsimin, asyinit, asyincr, asydecr
384 real(kind=rp),
intent(in),
optional :: move_limit
385 character(len=*),
intent(in),
optional :: bcknd, subsolver
386 character(len=256) :: log_msg
390 call this%scratch%init()
395 call this%xold1%init(n)
396 call this%xold2%init(n)
401 if (neko_bcknd_device .eq. 1)
then
402 call this%xold1%copy_from(device_to_host, sync = .false.)
403 call this%xold2%copy_from(device_to_host, sync = .true.)
406 call this%alpha%init(n)
407 call this%beta%init(n)
412 call this%low%init(n)
413 call this%upp%init(n)
414 call this%xmax%init(n)
415 call this%xmin%init(n)
418 call this%p0j%init(n)
419 call this%q0j%init(n)
420 call this%pij%init(m, n)
421 call this%qij%init(m, n)
426 call this%lambda%init(m)
429 call this%xsi%init(n)
430 call this%eta%init(n)
443 if (neko_bcknd_device .eq. 1)
then
444 call this%a%copy_from(host_to_device, sync = .false.)
445 call this%c%copy_from(host_to_device, sync = .false.)
446 call this%d%copy_from(host_to_device, sync = .false.)
447 call this%xmax%copy_from(host_to_device, sync = .false.)
448 call this%xmin%copy_from(host_to_device, sync = .true.)
452 this%residumax = huge(0.0_rp)
453 this%residunorm = huge(0.0_rp)
456 call mpi_allreduce(n, this%n_global, 1, mpi_integer, mpi_sum, neko_comm, &
463 if (.not.
present(max_iter)) this%max_iter = max_iter_default
464 if (.not.
present(epsimin))
then
465 this%epsimin = 1.0e-9_rp * sqrt(real(this%m + this%n_global, rp))
469 if (.not.
present(asyinit)) this%asyinit = asyinit_default
470 if (.not.
present(asyincr)) this%asyincr = asyincr_default
471 if (.not.
present(asydecr)) this%asydecr = asydecr_default
472 if (.not.
present(move_limit)) this%move_limit = move_limit_default
475 if (.not.
present(bcknd) .and. neko_bcknd_device .eq. 0)
then
477 else if (.not.
present(bcknd))
then
478 this%bcknd =
"device"
482 if (.not.
present(subsolver)) this%subsolver = subsolver_default
485 if (
present(max_iter)) this%max_iter = max_iter
486 if (
present(epsimin)) this%epsimin = epsimin
487 if (
present(asyinit)) this%asyinit = asyinit
488 if (
present(asyincr)) this%asyincr = asyincr
489 if (
present(asydecr)) this%asydecr = asydecr
490 if (
present(move_limit)) this%move_limit = move_limit
491 if (
present(bcknd)) this%bcknd = bcknd
492 if (
present(subsolver)) this%subsolver = subsolver
495 if (this%bcknd .ne.
"cpu" .and. this%bcknd .ne.
"device")
then
496 call neko_error(
"MMA: unknown backend '" // trim(this%bcknd) // &
497 "', expected 'cpu' or 'device'.")
501 if (.not. (this%epsimin .gt. 0.0_rp))
then
502 call neko_error(
"MMA: epsimin must be positive.")
505 call neko_log%section(
'MMA Parameters')
507 write(log_msg,
'(A10,1X,A)')
'backend ', trim(this%bcknd)
508 call neko_log%message(log_msg)
509 write(log_msg,
'(A10,1X,A)')
'subsolver ', trim(this%subsolver)
510 call neko_log%message(log_msg)
512 write(log_msg,
'(A10,1X,I0)')
'n ', this%n_global
513 call neko_log%message(log_msg)
514 write(log_msg,
'(A10,1X,I0)')
'm ', this%m
515 call neko_log%message(log_msg)
516 write(log_msg,
'(A10,1X,I0)')
'max_iter ', this%max_iter
517 call neko_log%message(log_msg)
519 write(log_msg,
'(A10,1X,E11.5)')
'epsimin ', this%epsimin
520 call neko_log%message(log_msg)
522 write(log_msg,
'(A10,1X,E11.5)')
'asyinit ', this%asyinit
523 call neko_log%message(log_msg)
524 write(log_msg,
'(A10,1X,E11.5)')
'asyincr ', this%asyincr
525 call neko_log%message(log_msg)
526 write(log_msg,
'(A10,1X,E11.5)')
'asydecr ', this%asydecr
527 call neko_log%message(log_msg)
528 write(log_msg,
'(A10,1X,E11.5)')
'a0 ', this%a0
529 call neko_log%message(log_msg)
530 write(log_msg,
'(A10,1X,E11.5)')
'movelimit ', this%move_limit
531 call neko_log%message(log_msg)
533 call neko_log%message(
'Parameters a')
535 write(log_msg,
'(3X,A,I2,A,E11.5)')
'a(', i,
') = ', this%a%x(i)
536 call neko_log%message(log_msg)
538 call neko_log%message(
'Parameters c')
540 write(log_msg,
'(3X,A,I2,A,E11.5)')
'c(', i,
') = ', this%c%x(i)
541 call neko_log%message(log_msg)
543 call neko_log%message(
'Parameters d')
545 write(log_msg,
'(3X,A,I2,A,E11.5)')
'd(', i,
') = ', this%d%x(i)
546 call neko_log%message(log_msg)
549 call neko_log%end_section()
552 this%is_initialized = .true.
553 end subroutine mma_init_from_components
559 subroutine mma_update_vector(this, iter, x, df0dx, fval, dfdx)
560 class(mma_t),
intent(inout) :: this
561 integer,
intent(in) :: iter
562 type(vector_t),
intent(inout) :: x
563 type(vector_t),
intent(inout) :: df0dx, fval
564 type(matrix_t),
intent(inout) :: dfdx
566 call profiler_start_region(
'FullMMAupdate')
568 select case (this%bcknd)
571 if (neko_bcknd_device .eq. 1)
then
572 call x%copy_from(device_to_host, sync = .false.)
573 call df0dx%copy_from(device_to_host, sync = .false.)
574 call fval%copy_from(device_to_host, sync = .false.)
575 call dfdx%copy_from(device_to_host, sync = .true.)
578 call mma_update_cpu(this, iter, x%x, df0dx%x, fval%x, dfdx%x)
580 if (neko_bcknd_device .eq. 1)
then
581 call x%copy_from(host_to_device, sync = .true.)
585 call mma_update_device(this, iter, x%x_d, df0dx%x_d, fval%x_d, dfdx%x_d)
587 call neko_error(
'mma_update_vector: Unknown backend: ' // &
590 call profiler_end_region(
'FullMMAupdate')
591 end subroutine mma_update_vector
594 subroutine mma_kkt_vector(this, x, df0dx, fval, dfdx)
595 class(mma_t),
intent(inout) :: this
596 type(vector_t),
intent(inout) :: x, df0dx, fval
597 type(matrix_t),
intent(inout) :: dfdx
600 select case (this%bcknd )
602 if (neko_bcknd_device .eq. 1)
then
603 call device_memcpy(x%x, x%x_d, this%n, device_to_host, &
605 call device_memcpy(df0dx%x, df0dx%x_d, this%n, device_to_host, &
607 call device_memcpy(fval%x, fval%x_d, this%m, device_to_host, &
609 call device_memcpy(dfdx%x, dfdx%x_d, this%m * this%n, device_to_host,&
613 call mma_kkt_cpu(this, x%x, df0dx%x, fval%x, dfdx%x)
615 call mma_kkt_device(this, x%x_d, df0dx%x_d, fval%x_d, dfdx%x_d)
617 call neko_error(
'mma_KKT_vector: Unknown backend: ' // &
620 end subroutine mma_kkt_vector
629 subroutine mma_save_checkpoint(this, filename, overwrite)
630 class(mma_t),
intent(inout) :: this
631 character(len=*),
intent(in) :: filename
632 logical,
intent(in),
optional :: overwrite
633 character(len=12) :: file_ext
636 call filename_suffix(filename, file_ext)
638 select case (trim(file_ext))
639 case (
'h5',
'hdf5',
'hf5')
640 call mma_save_checkpoint_hdf5(this, filename, overwrite)
642 call neko_error(
'mma_save_checkpoint: Unsupported file format: ' // &
646 end subroutine mma_save_checkpoint
651 subroutine mma_load_checkpoint(this, filename)
652 class(mma_t),
intent(inout) :: this
653 character(len=*),
intent(in) :: filename
654 character(len=12) :: file_ext
657 call filename_suffix(filename, file_ext)
659 select case (trim(file_ext))
660 case (
'h5',
'hdf5',
'hf5')
661 call mma_load_checkpoint_hdf5(this, filename)
663 call neko_error(
'mma_load_checkpoint: Unsupported file format: ' // &
666 end subroutine mma_load_checkpoint
672 pure function mma_get_n(this)
result(n)
673 class(mma_t),
intent(in) :: this
676 end function mma_get_n
679 pure function mma_get_m(this)
result(m)
680 class(mma_t),
intent(in) :: this
683 end function mma_get_m
686 pure function mma_get_residumax(this)
result(residumax)
687 class(mma_t),
intent(in) :: this
688 real(kind=rp) :: residumax
689 residumax = this%residumax
690 end function mma_get_residumax
693 pure function mma_get_residunorm(this)
result(residunorm)
694 class(mma_t),
intent(in) :: this
695 real(kind=rp) :: residunorm
696 residunorm = this%residunorm
697 end function mma_get_residunorm
700 pure function mma_get_max_iter(this)
result(max_iter_value)
701 class(mma_t),
intent(in) :: this
702 integer :: max_iter_value
703 max_iter_value = this%max_iter
704 end function mma_get_max_iter
707 pure function mma_get_backend_and_subsolver(this)
result(backend_subsolver)
708 class(mma_t),
intent(in) :: this
709 character(len=:),
allocatable :: backend_subsolver
710 character(len=:),
allocatable :: backend
712 if (neko_bcknd_cuda .eq. 1)
then
714 else if (neko_bcknd_hip .eq. 1)
then
716 else if (neko_bcknd_opencl .eq. 1)
then
722 backend_subsolver =
'backend:' // trim(backend) //
', subsolver:' // &
724 end function mma_get_backend_and_subsolver
736 subroutine mma_copy_from(this, direction, sync)
737 class(mma_t),
intent(inout) :: this
738 integer,
intent(in) :: direction
739 logical,
intent(in) :: sync
741 if (this%bcknd .ne.
"device")
return
743 call this%xold1%copy_from(direction, sync = .false.)
744 call this%xold2%copy_from(direction, sync = .false.)
745 call this%xmax%copy_from(direction, sync = .false.)
746 call this%xmin%copy_from(direction, sync = .false.)
748 call this%low%copy_from(direction, sync = .false.)
749 call this%upp%copy_from(direction, sync = .false.)
751 call this%a%copy_from(direction, sync = .false.)
752 call this%c%copy_from(direction, sync = .false.)
753 call this%d%copy_from(direction, sync = .false.)
754 call this%y%copy_from(direction, sync = .false.)
755 call this%s%copy_from(direction, sync = .false.)
757 call this%p0j%copy_from(direction, sync = .false.)
758 call this%q0j%copy_from(direction, sync = .false.)
759 call this%pij%copy_from(direction, sync = .false.)
760 call this%qij%copy_from(direction, sync = .false.)
761 call this%bi%copy_from(direction, sync = .false.)
763 call this%alpha%copy_from(direction, sync = .false.)
764 call this%beta%copy_from(direction, sync = .false.)
765 call this%lambda%copy_from(direction, sync = .false.)
766 call this%mu%copy_from(direction, sync = .false.)
767 call this%xsi%copy_from(direction, sync = .false.)
768 call this%eta%copy_from(direction, sync = sync)
770 end subroutine mma_copy_from
776 module subroutine mma_save_checkpoint_hdf5(object, filename, overwrite)
777 class(mma_t),
intent(inout) :: object
778 character(len=*),
intent(in) :: filename
779 logical,
intent(in),
optional :: overwrite
780 call neko_error(
'mma: HDF5 support not enabled rebuild with HAVE_HDF5')
781 end subroutine mma_save_checkpoint_hdf5
783 module subroutine mma_load_checkpoint_hdf5(object, filename)
784 class(mma_t),
intent(inout) :: object
785 character(len=*),
intent(in) :: filename
786 call neko_error(
'mma: HDF5 support not enabled rebuild with HAVE_HDF5')
787 end subroutine mma_load_checkpoint_hdf5