Move corpus related functions to corpus.py (google#63)

Author: Northbadge

compiler_opt/rl/compilation_runner.py

@@ -75,47 +75,6 @@ def _overwrite_trajectory_reward(sequence_example: tf.train.SequenceExample,
   return sequence_example
 
 
-def get_command_line_for_bundle(
-    cmd_file: str,
-    ir_file: str,
-    thinlto: Optional[str] = None,
-    additional_flags: Tuple[str, ...] = (),
-    delete_flags: Tuple[str, ...] = ()
-) -> List[str]:
-  """Cleans up base command line.
-
-  Remove certain unnecessary flags, and add the .bc file to compile and, if
-  given, the thinlto index.
-
-  Args:
-    cmd_file: Path to a .cmd file (from corpus).
-    ir_file: The path to the ir file to compile.
-    thinlto: The path to the thinlto index, or None.
-    additional_flags: Tuple of clang flags to add.
-    delete_flags: Tuple of clang flags to remove.
-
-  Returns:
-    The argument list to pass to the compiler process.
-  """
-  cmdline = []
-
-  with open(cmd_file, encoding='utf-8') as f:
-    option_iterator = iter(f.read().split('\0'))
-    option = next(option_iterator, None)
-    while option:
-      if any(option.startswith(flag) for flag in delete_flags):
-        if '=' not in option:
-          next(option_iterator, None)
-      else:
-        cmdline.append(option)
-      option = next(option_iterator, None)
-  cmdline.extend(['-x', 'ir', ir_file])
-  if thinlto:
-    cmdline.append('-fthinlto-index=' + thinlto)
-  cmdline.extend(additional_flags)
-  return cmdline
-
-
 class ProcessKilledError(Exception):
 
   def __init__(self):
@@ -286,23 +245,17 @@ def is_priority_method(cls, method_name: str) -> bool:
   def __init__(self,
                clang_path: Optional[str] = None,
                launcher_path: Optional[str] = None,
-               moving_average_decay_rate: float = 1,
-               additional_flags: Tuple[str, ...] = (),
-               delete_flags: Tuple[str, ...] = ()):
+               moving_average_decay_rate: float = 1):
     """Initialization of CompilationRunner class.
 
     Args:
       clang_path: path to the clang binary.
       launcher_path: path to the launcher binary.
       moving_average_decay_rate: moving average decay rate during training.
-      additional_flags: tuple of clang flags to add.
-      delete_flags: tuple of clang flags to remove.
     """
     self._clang_path = clang_path
     self._launcher_path = launcher_path
     self._moving_average_decay_rate = moving_average_decay_rate
-    self._additional_flags = additional_flags
-    self._delete_flags = delete_flags
     self._compilation_timeout = _COMPILATION_TIMEOUT.value
     self._cancellation_manager = WorkerCancellationManager()
 

compiler_opt/rl/compilation_runner_test.py

@@ -212,43 +212,6 @@ def test_exception_handling(self, mock_compile_fn):
           reward_stat=None)
     self.assertEqual(1, mock_compile_fn.call_count)
 
-  def test_command_line_file(self):
-    data = ['-cc1', '-foo', '-bar=baz']
-    argfile = self.create_tempfile(content='\0'.join(data))
-    self.assertEqual(
-        compilation_runner.get_command_line_for_bundle(argfile.full_path,
-                                                       'my_file.bc'),
-        ['-cc1', '-foo', '-bar=baz', '-x', 'ir', 'my_file.bc'])
-    self.assertEqual(
-        compilation_runner.get_command_line_for_bundle(argfile.full_path,
-                                                       'my_file.bc',
-                                                       'the_index.bc'),
-        [
-            '-cc1', '-foo', '-bar=baz', '-x', 'ir', 'my_file.bc',
-            '-fthinlto-index=the_index.bc'
-        ])
-
-  def test_command_line_correction(self):
-    delete_compilation_flags = ('-split-dwarf-file', '-split-dwarf-output',
-                                '-fthinlto-index', '-fprofile-sample-use',
-                                '-fprofile-remapping-file')
-    data = [
-        '-cc1', '-fthinlto-index=bad', '-split-dwarf-file', '/tmp/foo.dwo',
-        '-split-dwarf-output', 'somepath/some.dwo'
-    ]
-    argfile = self.create_tempfile(content='\0'.join(data))
-    self.assertEqual(
-        compilation_runner.get_command_line_for_bundle(
-            argfile.full_path, 'hi.bc', delete_flags=delete_compilation_flags),
-        ['-cc1', '-x', 'ir', 'hi.bc'])
-    self.assertEqual(
-        compilation_runner.get_command_line_for_bundle(
-            argfile.full_path,
-            'hi.bc',
-            'index.bc',
-            delete_flags=delete_compilation_flags),
-        ['-cc1', '-x', 'ir', 'hi.bc', '-fthinlto-index=index.bc'])
-
   def test_start_subprocess_output(self):
     ct = compilation_runner.WorkerCancellationManager()
     output = compilation_runner.start_cancellable_process(

compiler_opt/rl/corpus.py

@@ -15,11 +15,97 @@
 """ModuleSpec definition and utility command line parsing functions."""
 
 from dataclasses import dataclass
+from typing import Iterable, Optional, Tuple, List
+
+import os
+import tensorflow as tf
 
 
 @dataclass(frozen=True)
 class ModuleSpec:
   """Dataclass describing an input module and its compilation command options.
   """
   name: str
-  has_thinlto: bool = False
+  exec_cmd: Tuple[str, ...] = ()
+
+
+def build_modulespecs_from_datapath(
+    data_path: str,
+    additional_flags: Tuple[str, ...] = (),
+    delete_flags: Tuple[str, ...] = ()
+) -> List[ModuleSpec]:
+  module_paths: List[str] = _load_module_paths(data_path)
+
+  has_thinlto: bool = _has_thinlto_index(module_paths)
+
+  module_specs: List[ModuleSpec] = []
+
+  # This takes ~7s for 30k modules
+  for module_path in module_paths:
+    exec_cmd = _load_and_parse_command(
+        ir_file=module_path + '.bc',
+        cmd_file=(module_path + '.cmd'),
+        thinlto_file=(module_path + '.thinlto.bc') if has_thinlto else None,
+        additional_flags=additional_flags,
+        delete_flags=delete_flags)
+    module_specs.append(ModuleSpec(name=module_path, exec_cmd=tuple(exec_cmd)))
+
+  return module_specs
+
+
+def _has_thinlto_index(module_paths: Iterable[str]) -> bool:
+  return tf.io.gfile.exists(next(iter(module_paths)) + '.thinlto.bc')
+
+
+def _load_module_paths(data_path) -> List[str]:
+  module_paths_path = os.path.join(data_path, 'module_paths')
+  with open(module_paths_path, 'r', encoding='utf-8') as f:
+    ret = [os.path.join(data_path, name.rstrip('\n')) for name in f]
+    if len(ret) == 0:
+      raise ValueError(f'{module_paths_path} is empty.')
+    return ret
+
+
+def _load_and_parse_command(
+    ir_file: str,
+    cmd_file: str,
+    thinlto_file: Optional[str] = None,
+    additional_flags: Tuple[str, ...] = (),
+    delete_flags: Tuple[str, ...] = ()
+) -> List[str]:
+  """Cleans up base command line.
+
+  Remove certain unnecessary flags, and add the .bc file to compile and, if
+  given, the thinlto index.
+
+  Args:
+    cmd_file: Path to a .cmd file (from corpus).
+    ir_file: The path to the ir file to compile.
+    thinlto_file: The path to the thinlto index, or None.
+    additional_flags: Tuple of clang flags to add.
+    delete_flags: Tuple of clang flags to remove.
+
+  Returns:
+    The argument list to pass to the compiler process.
+  """
+  cmdline = []
+
+  with open(cmd_file, encoding='utf-8') as f:
+    option_iterator = iter(f.read().split('\0'))
+    option = next(option_iterator, None)
+    while option is not None:
+      if any(option.startswith(flag) for flag in delete_flags):
+        if '=' not in option:
+          next(option_iterator, None)
+      else:
+        cmdline.append(option)
+      option = next(option_iterator, None)
+  cmdline.extend(['-x', 'ir', ir_file])
+
+  if thinlto_file:
+    cmdline.extend(
+        [f'-fthinlto-index={thinlto_file}', '-mllvm', '-thinlto-assume-merged'])
+
+  cmdline.extend(additional_flags)
+
+  return cmdline