Add support for copying corpora to local disk

This patch adds support in regalloc_trace_worker for copying the corpus
to the local disk at initilization time. This is necessary for certain
filesystems that can only be interacted with through specific APIs,
precluding their direct use for applications like clang.

Reviewers: mtrofin

Reviewed By: mtrofin

Pull Request: #465

Author: boomanaiden154

compiler_opt/es/regalloc_trace/regalloc_trace_worker.py

@@ -27,15 +27,22 @@
 import concurrent.futures
 import tempfile
 import shutil
+from typing import Any
 
 import gin
+import tensorflow as tf
 
 from compiler_opt.rl import corpus
 from compiler_opt.distributed import worker
 from compiler_opt.rl import policy_saver
 from compiler_opt.es import policy_utils
 
 
+def _make_dirs_and_copy(old_file_path: str, new_file_path: str):
+  tf.io.gfile.makedirs(os.path.dirname(new_file_path))
+  tf.io.gfile.copy(old_file_path, new_file_path)
+
+
 @gin.configurable
 class RegallocTraceWorker(worker.Worker):
   """A worker that produces rewards for a given regalloc policy.
@@ -53,9 +60,55 @@ def _setup_base_policy(self):
     saver.save(self._tf_base_temp_dir)
     self._tf_base_policy_path = os.path.join(self._tf_base_temp_dir, "policy")
 
-  def __init__(self, *, gin_config: str, clang_path: str,
-               basic_block_trace_model_path: str, thread_count: int,
-               corpus_path: str):
+  def _copy_corpus(self, corpus_path: str,
+                   copy_corpus_locally_path: str | None) -> None:
+    """Makes a local copy of the corpus if requested.
+
+    This function makes a local copy of the corpus by copying the remote
+    corpus to a user-specified directory.
+
+    Args:
+      corpus_path: The path to the remote corpus.
+      copy_corpus_locally: The local path to copy the corpus to.
+    """
+    # We use the tensorflow APIs below rather than the standard Python file
+    # APIs for compatibility with more filesystems.
+
+    if tf.io.gfile.exists(copy_corpus_locally_path):
+      return
+
+    with tf.io.gfile.GFile(
+        os.path.join(corpus_path, "corpus_description.json"),
+        "r") as corpus_description_file:
+      corpus_description: dict[str, Any] = json.load(corpus_description_file)
+
+    file_extensions_to_copy = [".bc", ".cmd"]
+    if corpus_description["has_thinlto"]:
+      file_extensions_to_copy.append(".thinlto.bc")
+
+    copy_futures = []
+    with concurrent.futures.ThreadPoolExecutor(self._thread_count *
+                                               5) as copy_thread_pool:
+      for module in corpus_description["modules"]:
+        for extension in file_extensions_to_copy:
+          current_path = os.path.join(corpus_path, module + extension)
+          new_path = os.path.join(copy_corpus_locally_path, module + extension)
+          copy_futures.append(
+              copy_thread_pool.submit(_make_dirs_and_copy, current_path,
+                                      new_path))
+
+    for copy_future in copy_futures:
+      if copy_future.exception() is not None:
+        raise copy_future.exception()
+
+  def __init__(self,
+               *,
+               gin_config: str,
+               clang_path: str,
+               basic_block_trace_model_path: str,
+               thread_count: int,
+               corpus_path: str,
+               copy_corpus_locally_path: str | None = None):
     """Initializes the RegallocTraceWorker class.
 
     Args:
@@ -68,11 +121,19 @@ def __init__(self, *, gin_config: str, clang_path: str,
       thread_count: The number of threads to use for concurrent compilation
         and modelling.
       corpus_path: The path to the corpus that modules will be compiled from.
+      copy_corpus_locally_path: If set, specifies the path that the corpus
+        should be copied to before utilizing the modules for evaluation.
+        Setting this to None signifies that no copying is desired.
     """
     self._clang_path = clang_path
     self._basic_block_trace_model_path = basic_block_trace_model_path
     self._thread_count = thread_count
+    self._has_local_corpus = False
     self._corpus_path = corpus_path
+    if copy_corpus_locally_path is not None:
+      self._copy_corpus(corpus_path, copy_corpus_locally_path)
+      self._corpus_path = copy_corpus_locally_path
+      self._has_local_corpus = True
 
     gin.parse_config(gin_config)
     self._setup_base_policy()
@@ -83,6 +144,8 @@ def __init__(self, *, gin_config: str, clang_path: str,
   # have tempdirs wiped periodically.
   def __del__(self):
     shutil.rmtree(self._tf_base_temp_dir)
+    if self._has_local_corpus:
+      shutil.rmtree(self._corpus_path)
 
   def _compile_module(self, module_to_compile: corpus.ModuleSpec,
                       output_directory: str, tflite_policy_path: str | None):

compiler_opt/es/regalloc_trace/regalloc_trace_worker_test.py

@@ -15,6 +15,7 @@
 
 import os
 import json
+from pathlib import Path
 import stat
 import textwrap
 
@@ -26,15 +27,16 @@
 from compiler_opt.rl import corpus
 
 
-def _setup_corpus(corpus_dir: str) -> list[corpus.ModuleSpec]:
+def _setup_corpus(corpus_dir: str,
+                  has_thinlto: bool = False) -> list[corpus.ModuleSpec]:
   modules = [
-      corpus.ModuleSpec("module_a", 1, ("-fmodule-a",), True),
-      corpus.ModuleSpec("module_b", 1, ("-fmodule-b",), True)
+      corpus.ModuleSpec("module_a.o", 1, ("-fmodule-a",), True),
+      corpus.ModuleSpec("module_b.o", 1, ("-fmodule-b",), True)
   ]
 
   corpus_description = {
-      "has_thinlto": True,
-      "modules": [os.path.join(corpus_dir, module.name) for module in modules]
+      "has_thinlto": has_thinlto,
+      "modules": [module.name for module in modules]
   }
 
   with open(
@@ -43,6 +45,15 @@ def _setup_corpus(corpus_dir: str) -> list[corpus.ModuleSpec]:
       encoding="utf-8") as corpus_description_handle:
     json.dump(corpus_description, corpus_description_handle)
 
+  for module in ["module_a.o", "module_b.o"]:
+    extensions = [".cmd", ".bc"]
+    if has_thinlto:
+      extensions.append(".thinlto.bc")
+
+    for extension in extensions:
+      module_path = os.path.join(corpus_dir, module + extension)
+      Path(module_path).touch()
+
   return modules
 
 
@@ -151,3 +162,49 @@ def test_compile_corpus_and_evaluate_with_tflite(self):
     self.assertTrue(
         "-regalloc-enable-advisor=development" in clang_command_lines[1])
     self.assertTrue("-regalloc-model=" in clang_command_lines[1])
+
+  def test_copy_corpus_locally(self):
+    corpus_copy_base_dir = self.create_tempdir("corpus_copy")
+    corpus_copy_dir = os.path.join(corpus_copy_base_dir.full_path,
+                                   "corpus_copy")
+    corpus_dir = self.create_tempdir("corpus")
+    _ = _setup_corpus(corpus_dir.full_path)
+    worker = regalloc_trace_worker.RegallocTraceWorker(
+        gin_config="",
+        clang_path="/fake/path/to/clamg",
+        basic_block_trace_model_path="/fake/path/to/basic_block_trace_model",
+        thread_count=1,
+        corpus_path=corpus_dir.full_path,
+        copy_corpus_locally_path=corpus_copy_dir)
+
+    self.assertTrue(
+        os.path.exists(os.path.join(corpus_copy_dir, "module_a.o.bc")))
+    self.assertTrue(
+        os.path.exists(os.path.join(corpus_copy_dir, "module_a.o.cmd")))
+    self.assertTrue(
+        os.path.exists(os.path.join(corpus_copy_dir, "module_b.o.bc")))
+    self.assertTrue(
+        os.path.exists(os.path.join(corpus_copy_dir, "module_b.o.cmd")))
+
+    # Check that the worker cleans up after itself upon deletion.
+    del worker
+    self.assertFalse(os.path.exists(corpus_copy_dir))
+
+  def test_copy_corpus_locally_thinlto(self):
+    corpus_copy_base_dir = self.create_tempdir("corpus_copy")
+    corpus_copy_dir = os.path.join(corpus_copy_base_dir.full_path,
+                                   "corpus_copy")
+    corpus_dir = self.create_tempdir("corpus")
+    _ = _setup_corpus(corpus_dir.full_path, True)
+    _ = regalloc_trace_worker.RegallocTraceWorker(
+        gin_config="",
+        clang_path="/fake/path/to/clamg",
+        basic_block_trace_model_path="/fake/path/to/basic_block_trace_model",
+        thread_count=1,
+        corpus_path=corpus_dir.full_path,
+        copy_corpus_locally_path=corpus_copy_dir)
+
+    self.assertTrue(
+        os.path.exists(os.path.join(corpus_copy_dir, "module_a.o.thinlto.bc")))
+    self.assertTrue(
+        os.path.exists(os.path.join(corpus_copy_dir, "module_b.o.thinlto.bc")))