From f6766c169d1a7e133bd244ee6ac8ea44680b8d54 Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Sun, 1 Sep 2013 23:45:38 -0400
Subject: [PATCH] Add option to set number of threads in settings.xml and
 command line.

---
 src/global.F90               |  2 +-
 src/initialize.F90           | 17 +++++++++++++++++
 src/input_xml.F90            | 18 ++++++++++++++++++
 src/output.F90               |  1 +
 src/templates/settings.rnc   |  2 ++
 src/templates/settings_t.xml |  1 +
 6 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/src/global.F90 b/src/global.F90
index 1530f13a09..f96f59cf08 100644
--- a/src/global.F90
+++ b/src/global.F90
@@ -212,7 +212,7 @@ module global
   integer :: MPI_TALLYRESULT       ! MPI datatype for TallyResult
 
 #ifdef OPENMP
-  integer :: n_threads             ! number of OpenMP threads
+  integer :: n_threads = NONE      ! number of OpenMP threads
   integer :: thread_id             ! ID of a given thread
 #endif
 
diff --git a/src/initialize.F90 b/src/initialize.F90
index d82b1a7276..da2ace4d09 100644
--- a/src/initialize.F90
+++ b/src/initialize.F90
@@ -367,6 +367,23 @@ contains
         case ('-g', '-geometry-debug', '--geometry-debug')
           check_overlaps = .true.
 
+        case ('-t', '--threads')
+          ! Read number of threads
+          i = i + 1
+
+#ifdef OPENMP          
+          ! Read and set number of OpenMP threads
+          n_threads = str_to_int(argv(i))
+          if (n_threads < 1) then
+            message = "Invalid number of threads specified on command line."
+            call fatal_error()
+          end if
+          call omp_set_num_threads(n_threads)
+#else
+          message = "Ignoring number of threads specified on command line."
+          call warning()
+#endif
+
         case ('-?', '-help', '--help')
           call print_usage()
           stop
diff --git a/src/input_xml.F90 b/src/input_xml.F90
index 617788e472..c2d1a86d16 100644
--- a/src/input_xml.F90
+++ b/src/input_xml.F90
@@ -73,6 +73,7 @@ contains
     cross_sections_ = ''
     output_path_ = ''
     verbosity_ = 0
+    threads_ = NONE
     energy_grid_ = 'union'
     seed_ = 0_8
     source_ % file = ''
@@ -207,6 +208,23 @@ contains
     ! Verbosity
     if (verbosity_ > 0) verbosity = verbosity_
 
+    ! Number of OpenMP threads
+    if (threads_ /= NONE) then
+#ifdef OPENMP
+      if (n_threads == NONE) then
+        n_threads = threads_
+        if (n_threads < 1) then
+          message = "Invalid number of threads: " // to_str(n_threads)
+          call fatal_error()
+        end if
+        call omp_set_num_threads(n_threads)
+      end if
+#else
+      message = "Ignoring number of threads."
+      call warning()
+#endif
+    end if
+
     ! ==========================================================================
     ! EXTERNAL SOURCE
 
diff --git a/src/output.F90 b/src/output.F90
index 472cb5574b..ffffb745c5 100644
--- a/src/output.F90
+++ b/src/output.F90
@@ -170,6 +170,7 @@ contains
       write(OUTPUT_UNIT,*) '  -p, --plot             Run in plotting mode'
       write(OUTPUT_UNIT,*) '  -r, --restart          Restart a previous run from a state point'
       write(OUTPUT_UNIT,*) '                         or a particle restart file'
+      write(OUTPUT_UNIT,*) '  -t, --threads          Number of OpenMP threads'
       write(OUTPUT_UNIT,*) '  -v, --version          Show version information'
       write(OUTPUT_UNIT,*) '  -?, --help             Show this message'
     end if
diff --git a/src/templates/settings.rnc b/src/templates/settings.rnc
index 47dd690bcc..cd4ed4beb9 100644
--- a/src/templates/settings.rnc
+++ b/src/templates/settings.rnc
@@ -97,6 +97,8 @@ element settings {
 
   element survival_biasing { xsd:boolean }? &
 
+  element threads { xsd:positiveInteger }? &
+
   element trace { list { xsd:positiveInteger+ } }? &
 
   element verbosity { xsd:positiveInteger }? &
diff --git a/src/templates/settings_t.xml b/src/templates/settings_t.xml
index 3afa61745e..8f1cc7b00c 100644
--- a/src/templates/settings_t.xml
+++ b/src/templates/settings_t.xml
@@ -58,6 +58,7 @@
   <variable name="source_" tag="source" type="source_xml" />
   <variable name="state_point_" tag="state_point" type="statepoint_xml" dimension="1" />
   <variable name="survival_" tag="survival_biasing" type="word" length="5" />
+  <variable name="threads_" tag="threads" type="integer" />
   <variable name="trace_" tag="trace" type="integer-array" />
   <variable name="uniform_fs_" tag="uniform_fs" type="mesh_xml" dimension="1" />
   <variable name="verbosity_" tag="verbosity" type="integer" />