Switch file_paths to ModuleSpec (#57)

Preliminary patch to refactoring file_paths tuple usage completely.

Allows distributed thinlto corpora to be trained on

Author: Northbadge

compiler_opt/rl/compilation_runner.py

@@ -24,6 +24,7 @@
 from absl import flags
 from compiler_opt.distributed.worker import Worker, WorkerFuture
 from compiler_opt.rl import constant
+from compiler_opt.rl import corpus
 import tensorflow as tf
 
 _COMPILATION_TIMEOUT = flags.DEFINE_integer(
@@ -261,7 +262,7 @@ class CompilationRunnerStub(metaclass=abc.ABCMeta):
 
   @abc.abstractmethod
   def collect_data(
-      self, file_paths: Tuple[str, ...], tf_policy_path: str,
+      self, module_spec: corpus.ModuleSpec, tf_policy_path: str,
       reward_stat: Optional[Dict[str, RewardStat]]
   ) -> WorkerFuture[CompilationResult]:
     raise NotImplementedError()
@@ -313,12 +314,12 @@ def cancel_all_work(self):
     self._cancellation_manager.kill_all_processes()
 
   def collect_data(
-      self, file_paths: Tuple[str, ...], tf_policy_path: str,
+      self, module_spec: corpus.ModuleSpec, tf_policy_path: str,
       reward_stat: Optional[Dict[str, RewardStat]]) -> CompilationResult:
     """Collect data for the given IR file and policy.
 
     Args:
-      file_paths: path to files needed for inlining, Tuple of (.bc, .cmd).
+      module_spec: a ModuleSpec.
       tf_policy_path: path to the tensorflow policy.
       reward_stat: reward stat of this module, None if unknown.
       cancellation_token: a CancellationToken through which workers may be
@@ -336,7 +337,7 @@ def collect_data(
     """
     if reward_stat is None:
       default_result = self._compile_fn(
-          file_paths,
+          module_spec,
           tf_policy_path='',
           reward_only=bool(tf_policy_path),
           cancellation_manager=self._cancellation_manager)
@@ -346,7 +347,7 @@ def collect_data(
 
     if tf_policy_path:
       policy_result = self._compile_fn(
-          file_paths,
+          module_spec,
           tf_policy_path,
           reward_only=False,
           cancellation_manager=self._cancellation_manager)
@@ -362,7 +363,7 @@ def collect_data(
       if k not in reward_stat:
         raise ValueError(
             (f'Example {k} does not exist under default policy for '
-             'module {file_paths[0]}'))
+             f'module {module_spec.name}'))
       default_reward = reward_stat[k].default_reward
       moving_average_reward = reward_stat[k].moving_average_reward
       sequence_example = _overwrite_trajectory_reward(
@@ -384,13 +385,14 @@ def collect_data(
         keys=keys)
 
   def _compile_fn(
-      self, file_paths: Tuple[str, ...], tf_policy_path: str, reward_only: bool,
+      self, module_spec: corpus.ModuleSpec, tf_policy_path: str,
+      reward_only: bool,
       cancellation_manager: Optional[WorkerCancellationManager]
   ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
     """Compiles for the given IR file under the given policy.
 
     Args:
-      file_paths: path to files needed for compilation.
+      module_spec: a ModuleSpec.
       tf_policy_path: path to TF policy directory on local disk.
       reward_only: whether only return reward.
       cancellation_manager: a WorkerCancellationManager to handle early

compiler_opt/rl/compilation_runner_test.py

@@ -28,6 +28,7 @@
 
 from compiler_opt.rl import compilation_runner
 from compiler_opt.rl import constant
+from compiler_opt.rl import corpus
 
 _DEFAULT_FEATURE_VALUE = 12
 _POLICY_FEATURE_VALUE = 34
@@ -107,7 +108,7 @@ def test_policy(self, mock_compile_fn):
     runner = compilation_runner.CompilationRunner(
         moving_average_decay_rate=_MOVING_AVERAGE_DECAY_RATE)
     data = runner.collect_data(
-        file_paths=('bc', 'cmd'),
+        module_spec=corpus.ModuleSpec(name='dummy'),
         tf_policy_path='policy_path',
         reward_stat=None)
     self.assertEqual(2, mock_compile_fn.call_count)
@@ -138,7 +139,9 @@ def test_default(self, mock_compile_fn):
         moving_average_decay_rate=_MOVING_AVERAGE_DECAY_RATE)
 
     data = runner.collect_data(
-        file_paths=('bc', 'cmd'), tf_policy_path='', reward_stat=None)
+        module_spec=corpus.ModuleSpec(name='dummy'),
+        tf_policy_path='',
+        reward_stat=None)
     # One call when we ask for the default policy, because it can provide both
     # trace and default size.
     self.assertEqual(1, mock_compile_fn.call_count)
@@ -167,7 +170,7 @@ def test_given_default_size(self, mock_compile_fn):
         moving_average_decay_rate=_MOVING_AVERAGE_DECAY_RATE)
 
     data = runner.collect_data(
-        file_paths=('bc', 'cmd'),
+        module_spec=corpus.ModuleSpec(name='dummy'),
         tf_policy_path='policy_path',
         reward_stat={
             'default':
@@ -204,7 +207,7 @@ def test_exception_handling(self, mock_compile_fn):
 
     with self.assertRaisesRegex(subprocess.CalledProcessError, 'error'):
       _ = runner.collect_data(
-          file_paths=('bc', 'cmd'),
+          module_spec=corpus.ModuleSpec(name='dummy'),
           tf_policy_path='policy_path',
           reward_stat=None)
     self.assertEqual(1, mock_compile_fn.call_count)

compiler_opt/rl/corpus.py

@@ -0,0 +1,25 @@
+# coding=utf-8
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""ModuleSpec definition and utility command line parsing functions."""
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class ModuleSpec:
+  """Dataclass describing an input module and its compilation command options.
+  """
+  name: str
+  has_thinlto: bool = False

compiler_opt/rl/inlining/inlining_runner.py

@@ -23,6 +23,7 @@
 import tensorflow as tf
 
 from compiler_opt.rl import compilation_runner
+from compiler_opt.rl import corpus
 
 _DEFAULT_IDENTIFIER = 'default'
 
@@ -45,14 +46,14 @@ def __init__(self, llvm_size_path: str, *args, **kwargs):
     self._llvm_size_path = llvm_size_path
 
   def _compile_fn(
-      self, file_paths: Tuple[str, str], tf_policy_path: str, reward_only: bool,
-      cancellation_manager: Optional[
+      self, module_spec: corpus.ModuleSpec, tf_policy_path: str,
+      reward_only: bool, cancellation_manager: Optional[
           compilation_runner.WorkerCancellationManager]
   ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
     """Run inlining for the given IR file under the given policy.
 
     Args:
-      file_paths: path to files needed for inlining, Tuple of (.bc, .cmd).
+      module_spec: a ModuleSpec.
       tf_policy_path: path to TF policy direcoty on local disk.
       reward_only: whether only return native size.
       cancellation_manager: handler for early termination by killing any running
@@ -75,24 +76,22 @@ def _compile_fn(
     log_path = os.path.join(working_dir, 'log')
     output_native_path = os.path.join(working_dir, 'native')
 
-    input_ir_path, cmd_path = file_paths
-
     sequence_example = tf.train.SequenceExample()
     native_size = 0
     try:
       command_line = []
       if self._launcher_path:
         command_line.append(self._launcher_path)
-      command_line.extend([self._clang_path] +
-                          compilation_runner.get_command_line_for_bundle(
-                              cmd_path,
-                              input_ir_path,
-                              additional_flags=self._additional_flags,
-                              delete_flags=self._delete_flags) + [
-                                  '-mllvm', '-enable-ml-inliner=development',
-                                  '-mllvm', '-training-log=' +
-                                  log_path, '-o', output_native_path
-                              ])
+      command_line.extend(
+          [self._clang_path] + compilation_runner.get_command_line_for_bundle(
+              module_spec.name + '.cmd',
+              module_spec.name + '.bc', (module_spec.name + '.thinlto.bc'
+                                        ) if module_spec.has_thinlto else None,
+              additional_flags=self._additional_flags,
+              delete_flags=self._delete_flags) + [
+                  '-mllvm', '-enable-ml-inliner=development', '-mllvm',
+                  '-training-log=' + log_path, '-o', output_native_path
+              ])
       if tf_policy_path:
         command_line.extend(
             ['-mllvm', '-ml-inliner-model-under-training=' + tf_policy_path])

compiler_opt/rl/local_data_collector.py

@@ -25,6 +25,7 @@
 
 from compiler_opt.distributed import worker
 from compiler_opt.rl import compilation_runner
+from compiler_opt.rl import corpus
 from compiler_opt.rl import data_collector
 
 
@@ -33,7 +34,7 @@ class LocalDataCollector(data_collector.DataCollector):
 
   def __init__(
       self,
-      file_paths: Tuple[Tuple[str, ...], ...],
+      module_specs: List[corpus.ModuleSpec],
       num_modules: int,
       worker_pool: List[compilation_runner.CompilationRunnerStub],
       parser: Callable[[List[str]], Iterator[trajectory.Trajectory]],
@@ -43,7 +44,7 @@ def __init__(
     # TODO(mtrofin): type exit_checker_ctor when we get typing.Protocol support
     super().__init__()
 
-    self._file_paths = file_paths
+    self._module_specs = module_specs
     self._num_modules = num_modules
     self._parser = parser
     self._worker_pool = worker_pool
@@ -55,7 +56,7 @@ def __init__(
     # with the training phase - i.e. whatever happens between successive data
     # collection calls.
     self._reset_workers: concurrent.futures.Future = None
-    self._current_work: List[Tuple[Tuple[str, ...], worker.WorkerFuture]] = []
+    self._current_work: List[Tuple[corpus.ModuleSpec, worker.WorkerFuture]] = []
     self._pool = concurrent.futures.ThreadPoolExecutor()
 
   def close_pool(self):
@@ -77,14 +78,13 @@ def _join_pending_jobs(self):
                  time.time() - t1)
 
   def _schedule_jobs(
-      self, policy_path, sampled_file_paths
+      self, policy_path: str, sampled_modules: List[corpus.ModuleSpec]
   ) -> List[worker.WorkerFuture[compilation_runner.CompilationResult]]:
     # by now, all the pending work, which was signaled to cancel, must've
     # finished
     self._join_pending_jobs()
-    jobs = [(file_paths, policy_path,
-             self._reward_stat_map['-'.join(file_paths)])
-            for file_paths in sampled_file_paths]
+    jobs = [(module_spec, policy_path, self._reward_stat_map[module_spec.name])
+            for module_spec in sampled_modules]
 
     # Naive load balancing.
     ret = []
@@ -108,8 +108,8 @@ def collect_data(
       They will be reported using `tf.scalar.summary` by the trainer so these
       information is viewable in TensorBoard.
     """
-    sampled_file_paths = random.sample(self._file_paths, k=self._num_modules)
-    results = self._schedule_jobs(policy_path, sampled_file_paths)
+    sampled_modules = random.sample(self._module_specs, k=self._num_modules)
+    results = self._schedule_jobs(policy_path, sampled_modules)
 
     def wait_for_termination():
       early_exit = self._exit_checker_ctor(num_modules=self._num_modules)
@@ -121,12 +121,12 @@ def get_num_finished_work():
       return early_exit.wait(get_num_finished_work)
 
     wait_seconds = wait_for_termination()
-    self._current_work = list(zip(sampled_file_paths, results))
+    self._current_work = list(zip(sampled_modules, results))
     finished_work = [
-        (paths, res) for paths, res in self._current_work if res.done()
+        (spec, res) for spec, res in self._current_work if res.done()
     ]
-    successful_work = [(paths, res.result())
-                       for paths, res in finished_work
+    successful_work = [(spec, res.result())
+                       for spec, res in finished_work
                        if not worker.get_exception(res)]
     failures = len(finished_work) - len(successful_work)
 
@@ -149,10 +149,8 @@ def wrapup():
         itertools.chain.from_iterable(
             [res.serialized_sequence_examples for (_, res) in successful_work]))
     total_trajectory_length = sum(res.length for (_, res) in successful_work)
-    self._reward_stat_map.update({
-        '-'.join(file_paths): res.reward_stats
-        for (file_paths, res) in successful_work
-    })
+    self._reward_stat_map.update(
+        {spec.name: res.reward_stats for (spec, res) in successful_work})
 
     monitor_dict = {}
     monitor_dict['default'] = {

compiler_opt/rl/local_data_collector_test.py

@@ -25,6 +25,7 @@
 
 from compiler_opt.distributed.local.local_worker_manager import LocalWorkerPool
 from compiler_opt.rl import compilation_runner
+from compiler_opt.rl import corpus
 from compiler_opt.rl import data_collector
 from compiler_opt.rl import local_data_collector
 
@@ -46,8 +47,8 @@ def _get_sequence_example(feature_value):
   return text_format.Parse(sequence_example_text, tf.train.SequenceExample())
 
 
-def mock_collect_data(file_paths, tf_policy_dir, reward_stat):
-  assert file_paths == ('a', 'b')
+def mock_collect_data(module_spec, tf_policy_dir, reward_stat):
+  assert module_spec.name == 'dummy'
   assert tf_policy_dir == 'policy'
   assert reward_stat is None or reward_stat == {
       'default':
@@ -79,8 +80,8 @@ def mock_collect_data(file_paths, tf_policy_dir, reward_stat):
 class Sleeper(compilation_runner.CompilationRunner):
   """Test CompilationRunner that just sleeps."""
 
-  def collect_data(self, file_paths, tf_policy_path, reward_stat):
-    _ = file_paths, tf_policy_path, reward_stat
+  def collect_data(self, module_spec, tf_policy_path, reward_stat):
+    _ = module_spec, tf_policy_path, reward_stat
     compilation_runner.start_cancellable_process(['sleep', '3600s'], 3600,
                                                  self._cancellation_manager)
 
@@ -114,7 +115,7 @@ def _test_iterator_fn(data_list):
 
     with LocalWorkerPool(worker_class=MyRunner, count=4) as lwp:
       collector = local_data_collector.LocalDataCollector(
-          file_paths=tuple([('a', 'b')] * 100),
+          module_specs=[corpus.ModuleSpec(name='dummy')] * 100,
           num_modules=9,
           worker_pool=lwp,
           parser=create_test_iterator_fn(),
@@ -175,7 +176,7 @@ def wait(self, _):
 
     with LocalWorkerPool(worker_class=Sleeper, count=4) as lwp:
       collector = local_data_collector.LocalDataCollector(
-          file_paths=tuple([('a', 'b')] * 200),
+          module_specs=[corpus.ModuleSpec(name='dummy')] * 200,
           num_modules=4,
           worker_pool=lwp,
           parser=parser,