Refactor functionality of extract_ir.py into a library (#267)

This commit refactors most of the functional elements of extract_ir.py
into a library. This serves a couple of purposes. It makes some elements
of the code base a little bit cleaner (no more hacky flag tricks with
absl to get things working in unit tests) and can be seen as making some
things slightly more readable. In addition to this, it enables use of
these functions in other areas, particularly downstream projects within
the compiler_opt python package.

Author: boomanaiden154

compiler_opt/tools/extract_ir.py

@@ -32,19 +32,12 @@
 
 import json
 import multiprocessing
-import os
-import pathlib
-import re
-import shutil
-import subprocess
-
-from typing import Dict, List, Optional
 
 from absl import app
 from absl import flags
 from absl import logging
 
-from compiler_opt.rl import constant
+from compiler_opt.tools import extract_ir_lib
 
 flags.DEFINE_string(
     'input', None,
@@ -91,241 +84,6 @@
 FLAGS = flags.FLAGS
 
 
-# TODO(ml-compiler-opt): maybe we can also convert here the cmdline file,from a
-# \0 - separated list of strings, to a \n one.
-def should_include_module(cmdline: str, match_regexp: Optional[str]) -> bool:
-  """Determine if the module should be included."""
-  if match_regexp is None:
-    return True
-  lines = cmdline.split('\0')
-  return any(len(re.findall(match_regexp, l)) for l in lines)
-
-
-def get_thinlto_index(cmdline: str, basedir: str) -> Optional[str]:
-  opts = cmdline.split('\0')
-  for option in opts:
-    if option.startswith('-fthinlto-index'):
-      return os.path.join(basedir, option.split('=')[1])
-  return None
-
-
-class TrainingIRExtractor:
-  """IR and command line extraction from an object file."""
-
-  def __init__(self, obj_relative_path, output_base_dir, obj_base_dir=None):
-    """Set up a TrainingIRExtractor.
-
-    Args:
-      obj_relative_path: relative path to the input object file. It will be also
-        used to construct the absolute path of the output IR and cmd files, by
-        appending it to output_base_dir.
-      output_base_dir: the directory under which the output will be produced.
-      obj_base_dir: the base directory for all the input object files.
-    """
-    self._obj_relative_path = obj_relative_path
-    self._output_base_dir = output_base_dir
-    self._obj_base_dir = obj_base_dir if obj_base_dir is not None else ''
-
-  def obj_base_dir(self):
-    return self._obj_base_dir
-
-  def output_base_dir(self):
-    return self._output_base_dir
-
-  def relative_output_path(self):
-    return self._obj_relative_path
-
-  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))
-
-  def module_name(self):
-    return os.path.basename(self._obj_relative_path)
-
-  def cmd_file(self):
-    return os.path.join(self.dest_dir(), self.module_name() + '.cmd')
-
-  def bc_file(self):
-    return os.path.join(self.dest_dir(), self.module_name() + '.bc')
-
-  def thinlto_index_file(self):
-    return os.path.join(self.dest_dir(), self.module_name() + '.thinlto.bc')
-
-  def _get_extraction_cmd_command(self, llvm_objcopy_path):
-    """Call llvm_objcopy to extract the llvmcmd section in self._cmd_file."""
-    return [
-        llvm_objcopy_path,
-        '--dump-section=' + FLAGS.cmd_section_name + '=' + self.cmd_file(),
-        self.input_obj(), '/dev/null'
-    ]
-
-  def _get_extraction_bc_command(self, llvm_objcopy_path):
-    """Call llvm_objcopy to extract the llvmbc section in self._bc_file."""
-    return [
-        llvm_objcopy_path,
-        '--dump-section=' + FLAGS.bitcode_section_name + '=' + self.bc_file(),
-        self.input_obj(), '/dev/null'
-    ]
-
-  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())
-      return None
-    os.makedirs(self.dest_dir(), exist_ok=True)
-    try:
-      subprocess.run(
-          self._get_extraction_cmd_command(llvm_objcopy_path), check=True)
-      if cmd_filter is not None or is_thinlto:
-        with open(self.cmd_file(), encoding='utf-8') as f:
-          lines = f.readlines()
-        assert len(lines) == 1
-        cmdline = lines[0]
-        if not should_include_module(cmdline, cmd_filter):
-          logging.info(
-              'Excluding module %s because it does not match the filter',
-              self.input_obj())
-          os.remove(self.cmd_file())
-          return None
-        if is_thinlto:
-          index_file = get_thinlto_index(cmdline, self.obj_base_dir())
-          shutil.copy(index_file, self.thinlto_index_file())
-
-      subprocess.run(
-          self._get_extraction_bc_command(llvm_objcopy_path), check=True)
-    except subprocess.CalledProcessError as e:
-      # This may happen if  .o file was build from asm (.S source).
-      logging.warning('%s was not processed: %s', self.input_obj(), e)
-      return None
-    assert (os.path.exists(self.cmd_file()) and
-            os.path.exists(self.bc_file()) and
-            (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) -> Optional[TrainingIRExtractor]:
-  obj_base_dir = command['directory']
-  cmd = command['command']
-
-  cmd_parts = cmd.split()
-  try:
-    obj_index = cmd_parts.index('-o') + 1
-  except ValueError:
-    # This could happen if there are non-clang commands in compile_commands.json
-    logging.info('Command has no -o option: %s', cmd)
-    return None
-  obj_rel_path = cmd_parts[obj_index]
-  # TODO(mtrofin): is the obj_base_dir correct for thinlto index bc files?
-  return TrainingIRExtractor(
-      obj_relative_path=obj_rel_path,
-      output_base_dir=output_dir,
-      obj_base_dir=obj_base_dir)
-
-
-def load_from_compile_commands(json_array: List[Dict[str, str]],
-                               output_dir: str) -> List[TrainingIRExtractor]:
-  objs = [
-      convert_compile_command_to_objectfile(cmd, output_dir)
-      for cmd in json_array
-  ]
-  # Filter out None, in case there were non-clang commands in the .json
-  return [obj for obj in objs if obj is not None]
-
-
-def load_from_lld_params(params_array: List[str], obj_base_dir: str,
-                         output_dir: str) -> List[TrainingIRExtractor]:
-  """Create an ObjectFile array based on lld's parameters."""
-  # yank out -o and the output. After that, anything not starting with '-', and
-  # ending in a '.o', is an object file.
-  try:
-    minus_o_idx = params_array.index('-o')
-    del params_array[minus_o_idx:minus_o_idx + 2]
-    just_obj_paths = [
-        o for o in params_array if not o.startswith('-') and o.endswith('.o')
-    ]
-  except ValueError:
-    logging.info('This params file does not have an explicit -o option.')
-    just_obj_paths = params_array
-
-  def make_obj(obj_file: str) -> TrainingIRExtractor:
-    return TrainingIRExtractor(
-        obj_relative_path=obj_file,
-        output_base_dir=output_dir,
-        obj_base_dir=obj_base_dir)
-
-  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]:
-  return obj.extract(FLAGS.llvm_objcopy_path, FLAGS.cmd_filter,
-                     FLAGS.thinlto_build)
-
-
 def main(argv):
   if len(argv) > 1:
     raise app.UsageError('Too many command-line arguments.')
@@ -336,49 +94,34 @@ def main(argv):
   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)
+    objs = extract_ir_lib.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)
+      objs = extract_ir_lib.load_from_compile_commands(
+          json.load(f), FLAGS.output_dir)
   elif FLAGS.input_type == 'params':
     if not FLAGS.obj_base_dir:
       logging.info(
           '-obj_base_dir is unspecified, assuming current directory.'
           'If no objects are found, use this option to specify the root'
           'directory for the object file paths in the input file.')
     with open(FLAGS.input, encoding='utf-8') as f:
-      objs = load_from_lld_params([l.strip() for l in f.readlines()],
-                                  FLAGS.obj_base_dir, FLAGS.output_dir)
+      objs = extract_ir_lib.load_from_lld_params(
+          [l.strip() for l in f.readlines()], FLAGS.obj_base_dir,
+          FLAGS.output_dir)
   else:
     logging.error('Unknown input type: %s', FLAGS.input_type)
 
-  with multiprocessing.Pool(FLAGS.num_workers) as pool:
-    relative_output_paths = pool.map(extract_artifacts, objs)
-    pool.close()
-    pool.join()
-
-  # This comes first rather than later so global_command_override is at the top
-  # of the .json after being written
-  if FLAGS.thinlto_build == 'local':
-    corpus_description = {
-        'global_command_override': constant.UNSPECIFIED_OVERRIDE
-    }
-  else:
-    corpus_description = {}
-
-  corpus_description.update({
-      'has_thinlto': FLAGS.thinlto_build is not None,
-      'modules': [path for path in relative_output_paths if path is not None]
-  })
+  relative_output_paths = extract_ir_lib.run_extraction(
+      objs, FLAGS.num_workers, FLAGS.llvm_objcopy_path, FLAGS.cmd_filter,
+      FLAGS.thinlto_build, FLAGS.cmd_section_name, FLAGS.bitcode_section_name)
 
-  with open(
-      os.path.join(FLAGS.output_dir, 'corpus_description.json'),
-      'w',
-      encoding='utf-8') as f:
-    json.dump(corpus_description, f, indent=2)
+  extract_ir_lib.write_corpus_manifest(FLAGS.thinlto_build,
+                                       relative_output_paths, FLAGS.output_dir)
 
-    logging.info('Converted %d files out of %d',
-                 len(objs) - relative_output_paths.count(None), len(objs))
+  logging.info('Converted %d files out of %d',
+               len(objs) - relative_output_paths.count(None), len(objs))
 
 
 if __name__ == '__main__':