Add argument to append options to modelling invocations (#474)

This patch adds a new argument to regalloc_trace_worker that can be set
in a gin file for appending additional arguments to
basic_block_trace_model.

Author: boomanaiden154

compiler_opt/es/regalloc_trace/regalloc_trace_worker.py

@@ -125,6 +125,7 @@ def __init__(
       corpus_path: str,
       copy_corpus_locally_path: str | None = None,
       aux_file_replacement_flags: dict[str, str] | None = None,
+      extra_bb_trace_model_flags: list[str] | None = None,
   ):
     """Initializes the RegallocTraceWorker class.
 
@@ -146,10 +147,14 @@ def __init__(
         local to the worker. This is intended to be used in distributed
         training setups where training corpora and auxiliary files need to be
         copied locally before being compiled.
+      extra_bb_trace_model_flags: Extra flags to pass to the
+        basic_block_trace_model invocation.
     """
     self._clang_path = clang_path
     self._basic_block_trace_model_path = basic_block_trace_model_path
     self._thread_count = thread_count
+    self._extra_bb_trace_model_flags = ([] if extra_bb_trace_model_flags is None
+                                        else extra_bb_trace_model_flags)
 
     self._has_local_corpus = False
     self._corpus_path = corpus_path
@@ -257,6 +262,7 @@ def _evaluate_corpus(self, module_directory: str, function_index_path: str,
         f"--thread_count={self._thread_count}",
         f"--bb_trace_path={bb_trace_path}", "--model_type=mca"
     ]
+    command_vector.extend(self._extra_bb_trace_model_flags)
 
     output = subprocess.run(command_vector, capture_output=True, check=True)
 

compiler_opt/es/regalloc_trace/regalloc_trace_worker_test.py

@@ -252,3 +252,33 @@ def test_remote_corpus_replacement_flags(self):
         f"-fprofile-instr-use={copied_profile_path}" in clang_command_lines[0])
     self.assertTrue(
         f"-fprofile-instr-use={copied_profile_path}" in clang_command_lines[1])
+
+  def test_extra_bb_trace_flags(self):
+    corpus_dir = self.create_tempdir("corpus")
+    corpus_modules = _setup_corpus(corpus_dir.full_path)
+    fake_clang_binary = self.create_tempfile("fake_clang")
+    fake_clang_invocations = self.create_tempfile("fake_clang_invocations")
+    _create_test_binary(fake_clang_binary.full_path,
+                        fake_clang_invocations.full_path)
+    fake_bb_trace_model_binary = self.create_tempfile(
+        "fake_basic_block_trace_model")
+    fake_bb_trace_model_invocations = self.create_tempfile(
+        "fake_basic_block_trace_model_invocations")
+    _create_test_binary(fake_bb_trace_model_binary.full_path,
+                        fake_bb_trace_model_invocations.full_path)
+
+    worker = regalloc_trace_worker.RegallocTraceWorker(
+        gin_config="",
+        clang_path=fake_clang_binary.full_path,
+        basic_block_trace_model_path=fake_bb_trace_model_binary.full_path,
+        thread_count=1,
+        corpus_path=corpus_dir.full_path,
+        extra_bb_trace_model_flags=["--extra_flag"])
+    _ = worker.compile_corpus_and_evaluate(corpus_modules,
+                                           "function_index_path.pb",
+                                           "bb_trace_path.pb", None)
+
+    command_line = fake_bb_trace_model_invocations.read_text().split(
+        "\n")[0].split()
+
+    self.assertTrue("--extra_flag" in command_line)