Add Tao command to set number of OpenMP threads (#1357)

* Add Tao command to set number of OpenMP threads
* Set n_threads in 2 other places where globals are loaded

tao/code/tao_command.f90

@@ -692,7 +692,7 @@ subroutine tao_command (command_line, err_flag, err_is_fatal)
       'universe', 'curve', 'graph', 'beam_init', 'wave', 'plot', 'bmad_com', 'element', 'opti_de_param', &
       'csr_param', 'floor_plan', 'lat_layout', 'geodesic_lm', 'default', 'key', 'particle_start', &
       'plot_page', 'ran_state', 'symbolic_number', 'beam', 'beam_start', 'dynamic_aperture', &
-      'global', 'region', 'calculate', 'space_charge_com', 'ptc_com', 'tune', 'z_tune'], .true., switch, err_flag) 
+      'global', 'region', 'calculate', 'space_charge_com', 'ptc_com', 'tune', 'z_tune'], .true., switch, err_flag)
     if (err_flag) return
     set_word = switch
   enddo

tao/code/tao_de_optimizer.f90

@@ -18,6 +18,7 @@ subroutine tao_de_optimizer (abort)
 use tao_top10_mod, only: tao_var_write
 use opti_de_mod
 !!! use opti_de_openmp_mod  # And there is commented out code below.
+!!! use omp_lib, only: omp_get_max_threads
 
 implicit none
 
@@ -62,7 +63,7 @@ function merit_wrapper (var_vec, status, iter_count) result (merit)
 write (line, '(a, i0)') 'Differential evolution optimizer, population: ', population
 call out_io (s_blank$, r_name, line)
 
-!!! if (s%com%omp_n_threads == 1) then
+!!! if (omp_get_max_threads() == 1) then
   merit = opti_de (var_vec, s%global%n_opti_cycles, population, merit_wrapper, var_step, status)
 !!! else
 !!!  merit = opti_de_openmp (var_vec, s%global%n_opti_cycles, population, merit_wrapper, var_step, status)

tao/code/tao_init.f90

@@ -76,8 +76,8 @@ subroutine tao_init (err_flag)
 
 ! OpenMP info
 
-!$ s%com%omp_n_threads = omp_get_max_threads()
-!$ call out_io (s_important$, r_name, 'OpenMP active with number of threads: ' // int_str(s%com%omp_n_threads))
+!$ s%global%n_threads = omp_get_max_threads()
+!$ call out_io (s_important$, r_name, 'OpenMP active with number of threads: ' // int_str(s%global%n_threads))
 
 ! Open the init file.
 ! If the init file name is *not* the default (that is, it has been set by

tao/code/tao_init_mod.f90

@@ -23,6 +23,7 @@ subroutine tao_init_global (init_file)
 
 use opti_de_mod, only: opti_de_param
 use input_mod
+use tao_set_mod, only: tao_set_openmp_n_threads
 
 type (tao_global_struct) :: global
 
@@ -82,6 +83,8 @@ subroutine tao_init_global (init_file)
 
 close (iu)
 
+call tao_set_openmp_n_threads(s%global%n_threads)
+
 call end_bookkeeping()
 
 !-----------------------------------------------------------------------

tao/code/tao_parse_command_args.f90

@@ -14,6 +14,7 @@ subroutine tao_parse_command_args (error, cmd_line)
 
 use tao_interface, dummy => tao_parse_command_args
 use tao_command_mod, only: tao_cmd_split
+use tao_set_mod, only: tao_set_openmp_n_threads
 
 implicit none
 
@@ -114,6 +115,7 @@ subroutine tao_parse_command_args (error, cmd_line)
     s%init = tao_init_struct()
     s%com = tao_common0
     s%global = tao_global_struct()
+    call tao_set_openmp_n_threads(s%global%n_threads)
 
   case ('-command')
     call get_next_arg (arg0, s%init%command_arg, i_arg, n_arg, .true.)

tao/code/tao_pipe_cmd.f90

@@ -4479,6 +4479,7 @@ subroutine tao_pipe_cmd (input_str)
   nl=incr(nl); write (li(nl), amt) 'phase_units;ENUM;T;',                     trim(angle_units_name(s%global%phase_units))
   nl=incr(nl); write (li(nl), imt) 'bunch_to_plot;INT;T;',                    s%global%bunch_to_plot
   nl=incr(nl); write (li(nl), imt) 'random_seed;INT;T;',                      s%global%random_seed
+  nl=incr(nl); write (li(nl), imt) 'n_threads;INT;T;',                        s%global%n_threads
   nl=incr(nl); write (li(nl), imt) 'n_top10_merit;INT;T;',                    s%global%n_top10_merit
   nl=incr(nl); write (li(nl), imt) 'n_opti_loops;INT;T;',                     s%global%n_opti_loops
   nl=incr(nl); write (li(nl), imt) 'n_opti_cycles;INT;T;',                    s%global%n_opti_cycles

tao/code/tao_set_mod.f90

@@ -141,6 +141,49 @@ subroutine tao_set_z_tune_cmd (branch_str, q_str, delta_input)
 
 end subroutine tao_set_z_tune_cmd
 
+!-----------------------------------------------------------------------------
+!-----------------------------------------------------------------------------
+!-----------------------------------------------------------------------------
+!+
+! Subroutine tao_set_openmp_n_threads (n_threads)
+!
+! Routine to set OpenMP thread count.  Errors if OpenMP is not available.
+!
+! Input:
+!   n_threads      -- integer: Number of threads.
+!-
+
+subroutine tao_set_openmp_n_threads (n_threads)
+
+!$ use omp_lib, only: omp_get_max_threads, omp_set_num_threads
+
+implicit none
+
+integer old_n_threads, n_threads
+logical openmp_available
+
+character(*), parameter :: r_name = 'tao_set_openmp_n_threads'
+
+  openmp_available = .false.
+  !$ openmp_available = .true.
+
+  if (.not. openmp_available) then
+    if (n_threads > 1) then
+      call out_io (s_error$, r_name, 'Multithreading support with OpenMP is not available.')
+    endif
+    return
+  endif
+
+  !$ old_n_threads = omp_get_max_threads()
+  !$ call omp_set_num_threads(n_threads)
+  ! What OpenMP sets may differ from what we requested, so set it again here:
+  !$ s%global%n_threads = omp_get_max_threads()
+  !$ if (old_n_threads /= s%global%n_threads) then
+  !$   call out_io (s_important$, r_name, 'OpenMP active with number of threads: ' // int_str(s%global%n_threads))
+  !$ endif
+
+end subroutine tao_set_openmp_n_threads
+
 !-----------------------------------------------------------------------------
 !-----------------------------------------------------------------------------
 !-----------------------------------------------------------------------------
@@ -508,7 +551,8 @@ subroutine tao_set_global_cmd (who, value_str)
   val = quote(value_str)
 
 case ('n_opti_cycles', 'n_opti_loops', 'phase_units', 'bunch_to_plot', &
-      'random_seed', 'n_top10_merit', 'srdt_gen_n_slices', 'srdt_sxt_n_slices')
+      'random_seed', 'n_top10_merit', 'srdt_gen_n_slices', 'srdt_sxt_n_slices', &
+      'n_threads')
   call tao_evaluate_expression (value_str, 1, .false., set_val, err); if (err) return
   write (val, '(i0)', iostat = ios) nint(set_val(1))
 
@@ -547,6 +591,8 @@ subroutine tao_set_global_cmd (who, value_str)
 !
 
 select case (who)
+case ('n_threads')
+  call tao_set_openmp_n_threads(global%n_threads)
 case ('optimizer')
   if (all(global%optimizer /= tao_optimizer_name)) then
     call out_io (s_error$, r_name, 'BAD OPTIMIZER NAME: ' // global%optimizer)

tao/code/tao_show_this.f90

@@ -2133,6 +2133,7 @@ subroutine tao_show_this (what, result_id, lines, nl)
     nl=nl+1; write(lines(nl), lmt) '  %label_keys                    = ', s%global%label_keys
     nl=nl+1; write(lines(nl), lmt) '  %lattice_calc_on               = ', s%global%lattice_calc_on
     nl=nl+1; write(lines(nl), rmt) '  %max_plot_time                 = ', s%global%max_plot_time
+    nl=nl+1; write(lines(nl), imt) '  %n_threads                     = ', s%global%n_threads
     nl=nl+1; write(lines(nl), lmt) '  %only_limit_opt_vars           = ', s%global%only_limit_opt_vars
     nl=nl+1; write(lines(nl), lmt) '  %opt_match_auto_recalc         = ', s%global%opt_match_auto_recalc
     nl=nl+1; write(lines(nl), lmt) '  %opti_write_var_file           = ', s%global%opti_write_var_file

tao/code/tao_struct.f90

@@ -654,6 +654,7 @@ module tao_struct
   integer :: default_branch = 0                  ! Default lattice branch to work with.
   integer :: n_opti_cycles = 20                  ! Number of optimization cycles
   integer :: n_opti_loops = 1                    ! Number of optimization loops
+  integer :: n_threads = 1                       ! Number of OpenMP threads for parallel calculations.
   integer :: phase_units = radians$              ! Phase units on output.
   integer :: bunch_to_plot = 1                   ! Which bunch to plot
   integer :: random_seed = -1                    ! Use system clock by default
@@ -754,7 +755,6 @@ module tao_struct
   integer :: lev_loop = 0                       ! in do loop nest level
   integer :: n_err_messages_printed = 0         ! Used by tao_set_invalid to limit number of messages.
   integer :: n_universes = n_uni_init$   
-  integer :: omp_n_threads = 1                  ! Number of OpenMP threads.
   logical :: cmd_file_paused = .false.
   logical :: use_cmd_here  = .false.            ! Used for commands recalled from the cmd history stack
   logical :: cmd_from_cmd_file = .false.        ! was command from a command file?