Hopefully helpful modifications

compiler_opt/es/es_trainer_lib.py

@@ -76,10 +76,10 @@ def __init__(self, *, all_gin):
     self._template_dir = tempfile.mkdtemp()
     saver.save(self._template_dir)
 
-    self._corpus_dir = '/usr/local/google/home/aidengrossman/programming/test_traces/corpus_subset'
-    self._clang_path = '/usr/local/google/home/aidengrossman/programming/test_traces/clang'
-    self._trace_path = '/usr/local/google/home/aidengrossman/programming/test_traces/execution_trace.pb'
-    self._bb_trace_model_path = '/usr/local/google/home/aidengrossman/programming/test_traces/basic_block_trace_model'
+    self._corpus_dir = '/usr/local/google/home/aidengrossman/programming/opt_mlregalloc/corpus'
+    self._clang_path = '/usr/local/google/home/aidengrossman/programming/opt_mlregalloc/clang'
+    self._trace_path = '/usr/local/google/home/aidengrossman/programming/opt_mlregalloc/execution_trace.pb'
+    self._bb_trace_model_path = '/usr/local/google/home/aidengrossman/programming/opt_mlregalloc/basic_block_trace_model'
     self._models_for_test_path = '/usr/local/google/home/aidengrossman/programming/output_traces/'
 
   def es_compile(self, params: list[float], baseline_score: float) -> float:
@@ -98,15 +98,15 @@ def es_compile(self, params: list[float], baseline_score: float) -> float:
           os.path.join(tflitedir, policy_saver.OUTPUT_SIGNATURE))
 
       trace_data_collector.compile_corpus(self._corpus_dir, tempdir,
-                                          self._clang_path, tflitedir)
+                                          self._clang_path, tflitedir, single_threaded=True)
       score = trace_data_collector.evaluate_compiled_corpus(
           tempdir, self._trace_path, self._bb_trace_model_path)
 
       reward = compilation_runner._calculate_reward(score, baseline_score)
       print(reward)
 
       output_path = os.path.join(self._models_for_test_path, "model" + str(reward))
-      if reward > 0.2 and not os.path.exists(output_path):
+      if reward > 0 and not os.path.exists(output_path):
         shutil.copytree(tflitedir, output_path)
       return compilation_runner._calculate_reward(score, baseline_score)
 

compiler_opt/rl/trace_data_collector.py

@@ -17,42 +17,56 @@
 from absl import logging
 import os
 import subprocess
+import functools
 import json
 import pathlib
 import shutil
+import multiprocessing
 
 
-def compile_corpus(corpus_path, output_path, clang_path, tflite_dir=None):
+def compile_module(module_path, corpus_path, clang_path, tflite_dir, output_path):
+  module_full_input_path = os.path.join(corpus_path, module_path) + '.bc'
+  module_full_output_path = os.path.join(output_path, module_path) + '.bc.o'
+  module_command_full_path = os.path.join(corpus_path, module_path) + '.cmd'
+  pathlib.Path(os.path.dirname(module_full_output_path)).mkdir(
+      parents=True, exist_ok=True)
+  with open(module_command_full_path) as module_command_handle:
+    module_command_line = tuple(module_command_handle.read().replace(
+        r'{', r'{{').replace(r'}', r'}}').split('\0'))
+
+  command_vector = [clang_path]
+  command_vector.extend(module_command_line)
+  command_vector.extend(
+      [module_full_input_path, '-o', module_full_output_path])
+
+  if tflite_dir is not None:
+    command_vector.extend(['-mllvm', '-regalloc-enable-advisor=development'])
+    command_vector.extend(['-mllvm', '-regalloc-model=' + tflite_dir])
+
+  subprocess.run(command_vector, check=True)
+  logging.info(
+      f'Just finished compiling {module_full_output_path}'
+  )
+
+
+def compile_corpus(corpus_path, output_path, clang_path, tflite_dir=None, single_threaded=False):
   with open(
       os.path.join(corpus_path, 'corpus_description.json'),
       encoding='utf-8') as corpus_description_handle:
     corpus_description = json.load(corpus_description_handle)
 
+  to_compile = []
+
   for module_index, module_path in enumerate(corpus_description['modules']):
     # Compile each module.
-    module_full_input_path = os.path.join(corpus_path, module_path) + '.bc'
-    module_full_output_path = os.path.join(output_path, module_path) + '.bc.o'
-    pathlib.Path(os.path.dirname(module_full_output_path)).mkdir(
-        parents=True, exist_ok=True)
-
-    module_command_full_path = os.path.join(corpus_path, module_path) + '.cmd'
-    with open(module_command_full_path) as module_command_handle:
-      module_command_line = tuple(module_command_handle.read().replace(
-          r'{', r'{{').replace(r'}', r'}}').split('\0'))
-
-    command_vector = [clang_path]
-    command_vector.extend(module_command_line)
-    command_vector.extend(
-        [module_full_input_path, '-o', module_full_output_path])
-
-    if tflite_dir is not None:
-      command_vector.extend(['-mllvm', '-regalloc-enable-advisor=development'])
-      command_vector.extend(['-mllvm', '-regalloc-model=' + tflite_dir])
-
-    subprocess.run(command_vector, check=True)
-    logging.info(
-        f'Just finished compiling {module_full_output_path} ({module_index + 1}/{len(corpus_description["modules"])})'
-    )
+    to_compile.append(module_path)
+
+  if single_threaded:
+    for module_to_compile in to_compile:
+      compile_module(module_to_compile, corpus_path, clang_path, tflite_dir, output_path)
+  else:
+    with multiprocessing.Pool() as pool:
+      pool.map(functools.partial(compile_module, corpus_path=corpus_path, clang_path=clang_path, tflite_dir=tflite_dir, output_path=output_path), to_compile)
 
   shutil.copy(
       os.path.join(corpus_path, 'corpus_description.json'),
@@ -61,16 +75,17 @@ def compile_corpus(corpus_path, output_path, clang_path, tflite_dir=None):
 
 def evaluate_compiled_corpus(compiled_corpus_path, trace_path,
                              bb_trace_model_path):
+  corpus_description_path = os.path.join(compiled_corpus_path, 'corpus_description.json')
   command_vector = [
       bb_trace_model_path, '--bb_trace_path=' + trace_path,
-      '--corpus_path=' + compiled_corpus_path
+      '--corpus_path=' + corpus_description_path, '--cpu_name=skylake-avx512'
   ]
 
   process_return = subprocess.run(
       command_vector,
-      check=True,
       stdout=subprocess.PIPE,
       stderr=subprocess.STDOUT)
+
   output = process_return.stdout.decode('utf-8')
 
   return int(output)