Add support for extracting lld-thinlto corpora (#58)

Author: Northbadge

compiler_opt/tools/extract_ir.py

@@ -23,13 +23,17 @@
 The compilation is assumed to have been performed with clang, using
 -fembed-bitcode=all passed to cc1 (i.e. pass clang -Xclang=-fembed-bitcode=all)
 
-In a ThinLTO case, the compilation is assumed to have been performed specifying
--mllvm -lto-embed-bitcode=post-merge-pre-opt.
+In a distributed ThinLTO case, the compilation is assumed to have been performed
+specifying -mllvm -lto-embed-bitcode=post-merge-pre-opt.
+
+In a local ThinLTO case, the compilation is assumedto have been performed
+specifying -Wl,--save-temps=import -Wl,--thinlto-emit-index-files
 """
 
 import json
 import multiprocessing
 import os
+import pathlib
 import re
 import shutil
 import subprocess
@@ -59,11 +63,15 @@
     'Include only those modules with a command line matching this regexp. '
     'Setting it to None for not filtering. Note that the regexp is applied '
     'independently for each separate command line option. For example, ^-Oz$ '
-    'will match Oz - built binaries.')
-flags.DEFINE_bool(
-    'thinlto_build', False, 'Set if the build was ThinLTO, to '
-    'ensure index files are also copied. The build is assumed to have had'
-    '-mllvm -lto-embed-bitcode=post-merge-pre-opt passed to clang.')
+    'will match Oz - built binaries. Does not work with thinlto_build=lld.')
+flags.DEFINE_enum(
+    'thinlto_build', None, ['distributed', 'local'],
+    'Set if the build was performed with either \'distributed\' or '
+    '\'local\' ThinLTO. This ensures the thinlto.bc files are also copied. '
+    'The build is assumed to have had '
+    '-mllvm -lto-embed-bitcode=post-merge-pre-opt passed in the distributed '
+    'case, or -Wl,--save-temps=import and -Wl,--thinlto-emit-index-files '
+    'passed in the local case.')
 
 FLAGS = flags.FLAGS
 
@@ -118,6 +126,16 @@ def relative_output_path(self):
   def input_obj(self):
     return os.path.join(self.obj_base_dir(), self._obj_relative_path)
 
+  def lld_src_bc(self):
+    # .3.import.bc is the suffix attached to post-merge-pre-opt ('postimport')
+    # IR bitcode saved by lld. It is hardcoded into lld.
+    return os.path.join(self._obj_base_dir,
+                        self._obj_relative_path + '.3.import.bc')
+
+  def lld_src_thinlto(self):
+    return os.path.join(self._obj_base_dir,
+                        self._obj_relative_path + '.thinlto.bc')
+
   def dest_dir(self):
     return os.path.join(self.output_base_dir(),
                         os.path.dirname(self._obj_relative_path))
@@ -148,8 +166,8 @@ def _get_extraction_bc_command(self, llvm_objcopy_path):
         self.input_obj(), '/dev/null'
     ]
 
-  def extract(self, llvm_objcopy_path: str, cmd_filter: str,
-              is_thinlto: bool) -> Optional[str]:
+  def _extract_clang_artifacts(self, llvm_objcopy_path: str, cmd_filter: str,
+                               is_thinlto: bool) -> Optional[str]:
     """Run llvm-objcopy to extract the .bc and command line."""
     if not os.path.exists(self.input_obj()):
       logging.info('%s does not exist.', self.input_obj())
@@ -184,6 +202,36 @@ def extract(self, llvm_objcopy_path: str, cmd_filter: str,
             (not is_thinlto or os.path.exists(self.thinlto_index_file())))
     return self.relative_output_path()
 
+  def _extract_lld_artifacts(self) -> Optional[str]:
+    """Extract the .bc file with ThinLTO index from an lld ThinLTO invocation.
+    """
+    if not os.path.exists(self.lld_src_bc()):
+      logging.info('%s does not exist.', self.lld_src_bc())
+      return None
+    if not os.path.exists(self.lld_src_thinlto()):
+      logging.info('%s does not exist.', self.lld_src_thinlto())
+      return None
+    os.makedirs(self.dest_dir(), exist_ok=True)
+
+    # Copy over the files
+    shutil.copy(self.lld_src_bc(), self.bc_file())
+    shutil.copy(self.lld_src_thinlto(), self.thinlto_index_file())
+
+    assert os.path.exists(self.bc_file())
+    assert os.path.exists(self.thinlto_index_file())
+    return self._obj_relative_path
+
+  def extract(self,
+              llvm_objcopy_path: Optional[str] = None,
+              cmd_filter: Optional[str] = None,
+              thinlto_build: Optional[str] = None) -> Optional[str]:
+    if thinlto_build == 'local':
+      return self._extract_lld_artifacts()
+    return self._extract_clang_artifacts(
+        llvm_objcopy_path=llvm_objcopy_path,
+        cmd_filter=cmd_filter,
+        is_thinlto=thinlto_build == 'distributed')
+
 
 def convert_compile_command_to_objectfile(command: Dict[str, str],
                                           output_dir: str):
@@ -232,6 +280,24 @@ def make_obj(obj_file: str) -> TrainingIRExtractor:
   return [make_obj(obj_file) for obj_file in just_obj_paths]
 
 
+def load_for_lld_thinlto(obj_base_dir: str,
+                         output_dir: str) -> List[TrainingIRExtractor]:
+  # .3.import.bc is the suffix attached to post-merge-pre-opt ('postimport')
+  # IR bitcode saved by lld. It is hardcoded into lld. ThinLTO index files
+  # are also emitted next to the postimport bitcode, with the suffix
+  # .thinlto.bc instead
+  paths = [str(p) for p in pathlib.Path(obj_base_dir).glob('**/*.3.import.bc')]
+
+  def make_spec(obj_file: str):
+    return TrainingIRExtractor(
+        # Cut away .3.import.bc
+        obj_relative_path=os.path.relpath(obj_file, start=obj_base_dir)[:-12],
+        output_base_dir=output_dir,
+        obj_base_dir=obj_base_dir)
+
+  return [make_spec(path) for path in paths]
+
+
 # This is here just for readability, lint complains if the pooling expression is
 # over 3 lines; and it needs to be a non-local so it may be pickled.
 def extract_artifacts(obj: TrainingIRExtractor) -> Optional[str]:
@@ -242,9 +308,14 @@ def extract_artifacts(obj: TrainingIRExtractor) -> Optional[str]:
 def main(argv):
   if len(argv) > 1:
     raise app.UsageError('Too many command-line arguments.')
-  flags.mark_flags_as_required(['output_dir', 'input'])
+  flags.mark_flags_as_required(['output_dir'])
+
   objs = []
-  if FLAGS.input_type == 'json':
+  if FLAGS.input is None:
+    if FLAGS.thinlto_build != 'local':
+      raise ValueError('--input or --thinlto_build=local must be provided')
+    objs = load_for_lld_thinlto(FLAGS.obj_base_dir, FLAGS.output_dir)
+  elif FLAGS.input_type == 'json':
     with open(FLAGS.input, encoding='utf-8') as f:
       objs = load_from_compile_commands(json.load(f), FLAGS.output_dir)
   elif FLAGS.input_type == 'params':

compiler_opt/tools/extract_ir_test.py

@@ -15,6 +15,7 @@
 """Tests for compiler_opt.tools.extract_ir."""
 
 # pylint: disable=protected-access
+import os.path
 
 from absl import flags
 from absl.testing import absltest
@@ -116,6 +117,63 @@ def test_lld_params(self):
                      '/tmp/out/lib/obj1.o.thinlto.bc')
     self.assertEqual(obj[1].input_obj(), '/some/path/lib/dir/obj2.o')
 
+  def test_lld_thinlto_discovery(self):
+    tempdir = self.create_tempdir()
+    tempdir.create_file(file_path='1.3.import.bc')
+    tempdir.create_file(file_path='2.3.import.bc')
+    tempdir.create_file(file_path='3.3.import.bc')
+    tempdir.create_file(file_path='1.thinlto.bc')
+    tempdir.create_file(file_path='2.thinlto.bc')
+    tempdir.create_file(file_path='3.thinlto.bc')
+    outdir = self.create_tempdir()
+    obj = extract_ir.load_for_lld_thinlto(tempdir.full_path, outdir.full_path)
+    self.assertLen(obj, 3)
+    for i, o in enumerate(sorted(obj, key=lambda x: x._obj_relative_path)):
+      self.assertEqual(o._obj_relative_path, f'{i + 1:d}')
+      self.assertEqual(o._obj_base_dir, tempdir.full_path)
+      self.assertEqual(o._output_base_dir, outdir.full_path)
+
+  def test_lld_thinlto_discovery_nested(self):
+    outer = self.create_tempdir()
+    tempdir = outer.mkdir(dir_path='nest')
+    tempdir.create_file(file_path='1.3.import.bc')
+    tempdir.create_file(file_path='2.3.import.bc')
+    tempdir.create_file(file_path='3.3.import.bc')
+    tempdir.create_file(file_path='1.thinlto.bc')
+    tempdir.create_file(file_path='2.thinlto.bc')
+    tempdir.create_file(file_path='3.thinlto.bc')
+    outdir = self.create_tempdir()
+    obj = extract_ir.load_for_lld_thinlto(outer.full_path, outdir.full_path)
+    self.assertLen(obj, 3)
+    for i, o in enumerate(sorted(obj, key=lambda x: x._obj_relative_path)):
+      self.assertEqual(o._obj_relative_path, f'nest/{i + 1:d}')
+      self.assertEqual(o._obj_base_dir, outer.full_path)
+      self.assertEqual(o._output_base_dir, outdir.full_path)
+
+  def test_lld_thinlto_extraction(self):
+    outer = self.create_tempdir()
+    tempdir = outer.mkdir(dir_path='nest')
+    tempdir.create_file(file_path='1.3.import.bc')
+    tempdir.create_file(file_path='2.3.import.bc')
+    tempdir.create_file(file_path='3.3.import.bc')
+    tempdir.create_file(file_path='1.thinlto.bc')
+    tempdir.create_file(file_path='2.thinlto.bc')
+    tempdir.create_file(file_path='3.thinlto.bc')
+    outdir = self.create_tempdir()
+    obj = extract_ir.load_for_lld_thinlto(outer.full_path, outdir.full_path)
+    for i, o in enumerate(sorted(obj, key=lambda x: x._obj_relative_path)):
+      mod_path = o.extract(thinlto_build='local')
+      self.assertEqual(mod_path, f'nest/{i + 1:d}')
+    self.assertTrue(os.path.exists(os.path.join(outdir.full_path, 'nest/1.bc')))
+    self.assertTrue(os.path.exists(os.path.join(outdir.full_path, 'nest/2.bc')))
+    self.assertTrue(os.path.exists(os.path.join(outdir.full_path, 'nest/3.bc')))
+    self.assertTrue(
+        os.path.exists(os.path.join(outdir.full_path, 'nest/1.thinlto.bc')))
+    self.assertTrue(
+        os.path.exists(os.path.join(outdir.full_path, 'nest/2.thinlto.bc')))
+    self.assertTrue(
+        os.path.exists(os.path.join(outdir.full_path, 'nest/3.thinlto.bc')))
+
   def test_filtering(self):
     cmdline = '-cc1\0x/y/foobar.cpp\0-Oz\0-Ifoo\0-o\0bin/out.o'
     self.assertTrue(extract_ir.should_include_module(cmdline, None))