Refactor add/delete flags options to problem_configuration (#84)

This patch moves the add/delete flags from a gin config option to the
problem configuration python object and makes generate_default_trace
grab the flags from there.

Author: boomanaiden154

compiler_opt/rl/inlining/gin_configs/ppo_nn_agent.gin

@@ -20,18 +20,6 @@ train_eval.deploy_policy_name='saved_collect_policy'
 train_eval.use_random_network_distillation=False
 train_eval.moving_average_decay_rate=0.8
 
-# 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
 RandomNetworkDistillation.learning_rate = 1e-4

compiler_opt/rl/problem_configuration.py

@@ -99,3 +99,19 @@ def get_nonnormalized_features(self) -> Iterable[str]:
   @abc.abstractmethod
   def get_runner_type(self) -> 'type[compilation_runner.CompilationRunner]':
     raise NotImplementedError
+
+  # 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:
+  # return ['-fprofile-sample-use=/path/to/gwp.afdo',
+  #  '-fprofile-remapping-file=/path/to/prof_remap.txt']
+  def flags_to_add(self) -> Tuple[str, ...]:
+    return ()
+
+  # List of flags to remove from clang compilation command. The flag names
+  # should match the actual flags provided to clang.'
+  def flags_to_delete(self) -> Tuple[str, ...]:
+    return ('-split-dwarf-file', '-split-dwarf-output', '-fthinlto-index',
+            '-fprofile-sample-use', '-fprofile-remapping-file')

compiler_opt/rl/regalloc/gin_configs/ppo_nn_agent.gin

@@ -21,18 +21,6 @@ train_eval.deploy_policy_name='saved_collect_policy'
 train_eval.moving_average_decay_rate=0.8
 train_eval.use_random_network_distillation=False
 
-# 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
 RandomNetworkDistillation.learning_rate = 1e-4

compiler_opt/rl/train_locally.py

@@ -68,9 +68,7 @@ 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,
-               additional_compilation_flags=(),
-               delete_compilation_flags=()):
+               moving_average_decay_rate=1):
   """Train for LLVM inliner."""
   root_dir = FLAGS.root_dir
   problem_config = registry.get_configuration()
@@ -102,7 +100,8 @@ def train_eval(agent_name=constant.AgentName.PPO,
 
   logging.info('Loading module specs from corpus.')
   module_specs = corpus.build_modulespecs_from_datapath(
-      FLAGS.data_path, additional_compilation_flags, delete_compilation_flags)
+      FLAGS.data_path, problem_config.flags_to_add(),
+      problem_config.flags_to_delete())
   logging.info('Done loading module specs from corpus.')
 
   dataset_fn = data_reader.create_sequence_example_dataset_fn(

compiler_opt/tools/generate_default_trace.py

@@ -136,12 +136,13 @@ def main(_):
       _GIN_FILES.value, bindings=_GIN_BINDINGS.value, skip_unknown=False)
   logging.info(gin.config_str())
 
+  config = registry.get_configuration()
+
   logging.info('Loading module specs from corpus.')
   module_specs = corpus.build_modulespecs_from_datapath(
       _DATA_PATH.value,
-      delete_flags=('-split-dwarf-file', '-split-dwarf-output',
-                    '-fthinlto-index', '-fprofile-sample-use',
-                    '-fprofile-remapping-file'))
+      additional_flags=config.flags_to_add(),
+      delete_flags=config.flags_to_delete())
   logging.info('Done loading module specs from corpus.')
 
   if _MODULE_FILTER.value:

compiler_opt/tools/generate_default_trace_test.py

@@ -60,9 +60,16 @@ def collect_data(self, module_spec, tf_policy_path, reward_stat):
 
 
 class GenerateDefaultTraceTest(absltest.TestCase):
+  def setUp(self):
+    with gin.unlock_config():
+      gin.parse_config_files_and_bindings(
+          config_files=['compiler_opt/rl/inlining/gin_configs/common.gin'],
+          bindings=None)
+    return super().setUp()
 
   @mock.patch('compiler_opt.tools.generate_default_trace.get_runner')
   def test_api(self, mock_get_runner):
+
     tmp_dir = self.create_tempdir()
     module_names = ['a', 'b', 'c', 'd']
 
@@ -92,10 +99,6 @@ def test_api(self, mock_get_runner):
       generate_default_trace.main(None)
 
   def test_get_runner(self):
-    with gin.unlock_config():
-      gin.parse_config_files_and_bindings(
-          config_files=['compiler_opt/rl/inlining/gin_configs/common.gin'],
-          bindings=None)
     runner = generate_default_trace.get_runner()
     self.assertIsInstance(runner, compilation_runner.CompilationRunner)