Modify OMP env variables setting

Author: Alexsandruss

runner.py

@@ -97,6 +97,9 @@ def is_ht_enabled():
 with open(args.config.name, 'r') as config_file:
     config = json.load(config_file)
 
+if 'omp_env' not in config.keys():
+    config['omp_env'] = []
+
 # make directory for data if it doesn't exist
 os.makedirs('data', exist_ok=True)
 
@@ -168,6 +171,11 @@ def is_ht_enabled():
 cpu_count = multiprocessing.cpu_count()
 omp_num_threads = str(cpu_count // 2) if is_ht_enabled() else str(cpu_count)
 
+omp_env = {
+    'OMP_PLACES': f'{{0}}:{cpu_count}:1',
+    'OMP_NUM_THREADS': omp_num_threads
+}
+
 # get parameters that are common for all cases
 common_params = config['common']
 for params_set in config['cases']:
@@ -261,8 +269,8 @@ class GenerationArgs:
         for lib in libs:
             env = os.environ.copy()
             if lib == 'xgboost':
-                env['OMP_NUM_THREADS'] = omp_num_threads
-
+                for var in config['omp_env']:
+                    env[var] = omp_env[var]
             for i, case in enumerate(cases):
                 command = f'python {lib}/{algorithm}.py --batch {batch} ' \
                           + f'--arch {hostname} --header --output-format ' \

xgboost/gbt.py

@@ -9,6 +9,7 @@
 )
 import numpy as np
 import xgboost as xgb
+import os
 
 
 def convert_probs_to_classes(y_prob):
@@ -97,6 +98,9 @@ def convert_xgb_predictions(y_pred, objective):
 if params.threads != -1:
     xgb_params.update({'nthread': params.threads})
 
+if 'OMP_NUM_THREADS' in os.environ.keys():
+    xgb_params['nthread'] = int(os.environ['OMP_NUM_THREADS'])
+
 columns = ('batch', 'arch', 'prefix', 'function', 'threads', 'dtype', 'size',
            'num_trees')