Refactors "to be deleted" and "to be added" compilation flags and enables training with AFDO profile reingestion.

Author: kshiteejm

compiler_opt/rl/compilation_runner.py

@@ -75,9 +75,12 @@ 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) -> List[str]:
+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
@@ -87,21 +90,19 @@ def get_command_line_for_bundle(cmd_file: str,
     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 = []
-  flags_to_remove = [
-      '-split-dwarf-file', '-split-dwarf-output', '-fthinlto-index',
-      '-fprofile-sample-use'
-  ]
 
   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 flags_to_remove):
+      if any([option.startswith(flag) for flag in delete_flags]):
         if '=' not in option:
           next(option_iterator, None)
       else:
@@ -110,6 +111,7 @@ def get_command_line_for_bundle(cmd_file: str,
   cmdline.extend(['-x', 'ir', ir_file])
   if thinlto:
     cmdline.append('-fthinlto-index=' + thinlto)
+  cmdline.extend(additional_flags)
   return cmdline
 
 
@@ -281,17 +283,23 @@ def _get_pool():
   def __init__(self,
                clang_path: Optional[str] = None,
                launcher_path: Optional[str] = None,
-               moving_average_decay_rate: float = 1):
+               moving_average_decay_rate: float = 1,
+               additional_flags: Tuple[str, ...] = (),
+               delete_flags: Tuple[str, ...] = ()):
     """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
 
   def _get_cancellation_manager(

compiler_opt/rl/compilation_runner_test.py

@@ -224,18 +224,24 @@ def test_command_line_file(self):
         ])
 
   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'),
+        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'),
+        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):

compiler_opt/rl/inlining/gin_configs/ppo_nn_agent.gin

@@ -19,10 +19,17 @@ train_eval.deploy_policy_name='saved_collect_policy'
 train_eval.use_random_network_distillation=False
 train_eval.moving_average_decay_rate=0.8
 
-#######################################
-# Turn on if using train_with_rpc.py
-# train_eval.additional_compilation_flags=()
-#######################################
+# TODO(b/233935329): The following clang flags need to be tied to a corpus
+# rather than to a training tool invocation.
+
+# List of flags to add to clang compilation command. The flag names should
+# match the actual flags provided to clang. An example for AFDO reinjection:
+# train_eval.flags_to_add=('-fprofile-sample-use=/path/to/gwp.afdo','-fprofile-remapping-file=/path/to/prof_remap.txt')
+train_eval.additional_compilation_flags=()
+
+# List of flags to remove from clang compilation command. The flag names 
+# should match the actual flags provided to clang.'
+train_eval.delete_compilation_flags=('-split-dwarf-file','-split-dwarf-output','-fthinlto-index','-fprofile-sample-use','-fprofile-remapping-file')
 
 # RandomNetworkDistillation configs, off if train_eval.use_random_network_distillation=False.
 RandomNetworkDistillation.encoding_network = @encoding_network.EncodingNetwork

compiler_opt/rl/inlining/inlining_runner.py

@@ -85,7 +85,10 @@ def _compile_fn(
         command_line.append(self._launcher_path)
       command_line.extend([self._clang_path] +
                           compilation_runner.get_command_line_for_bundle(
-                              cmd_path, input_ir_path) + [
+                              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

compiler_opt/rl/regalloc/gin_configs/ppo_nn_agent.gin

@@ -20,10 +20,17 @@ train_eval.deploy_policy_name='saved_collect_policy'
 train_eval.moving_average_decay_rate=0.8
 train_eval.use_random_network_distillation=False
 
-#######################################
-# Turn on if using train_with_rpc.py
-# train_eval.additional_compilation_flags=()
-#######################################
+# TODO(b/233935329): The following clang flags need to be tied to a corpus
+# rather than to a training tool invocation.
+
+# List of flags to add to clang compilation command. The flag names should
+# match the actual flags provided to clang. An example for AFDO reinjection:
+# train_eval.additional_flags=('-fprofile-sample-use=/path/to/gwp.afdo','-fprofile-remapping-file=/path/to/prof_remap.txt')
+train_eval.additional_compilation_flags=()
+
+# List of flags to remove from clang compilation command. The flag names 
+# should match the actual flags provided to clang.'
+train_eval.delete_compilation_flags=('-split-dwarf-file','-split-dwarf-output','-fthinlto-index','-fprofile-sample-use','-fprofile-remapping-file')
 
 # RandomNetworkDistillation configs, off if train_eval.use_random_network_distillation=False.
 RandomNetworkDistillation.encoding_network = @regalloc_network.RegAllocRNDEncodingNetwork

compiler_opt/rl/regalloc/regalloc_runner.py

@@ -89,7 +89,8 @@ def _compile_fn(
         command_line.append(self._launcher_path)
       command_line.extend([self._clang_path] +
                           compilation_runner.get_command_line_for_bundle(
-                              cmd_path, input_ir_path, thinlto_index_path) + [
+                              cmd_path, input_ir_path, thinlto_index_path,
+                              self._additional_flags, self._delete_flags) + [
                                   '-mllvm', '-thinlto-assume-merged', '-mllvm',
                                   '-regalloc-enable-advisor=development',
                                   '-mllvm', '-regalloc-training-log=' +

compiler_opt/rl/train_locally.py

@@ -67,7 +67,9 @@ def train_eval(agent_name=constant.AgentName.PPO,
                train_sequence_length=1,
                deploy_policy_name='saved_policy',
                use_random_network_distillation=False,
-               moving_average_decay_rate=1):
+               moving_average_decay_rate=1,
+               additional_compilation_flags=(),
+               delete_compilation_flags=()):
   """Train for LLVM inliner."""
   root_dir = FLAGS.root_dir
   problem_config = registry.get_configuration()
@@ -112,7 +114,9 @@ def train_eval(agent_name=constant.AgentName.PPO,
                     for path in module_paths]
 
   runner = problem_config.get_runner_type()(
-      moving_average_decay_rate=moving_average_decay_rate)
+      moving_average_decay_rate=moving_average_decay_rate,
+      additional_flags=additional_compilation_flags,
+      delete_flags=delete_compilation_flags)
 
   dataset_fn = data_reader.create_sequence_example_dataset_fn(
       agent_name=agent_name,