Add directory input to extract_ir.py

This patch adds in a directory input to extract_ir.py where the script
does a glob for *.o files inside the specified directory. This is useful
in cases where it is easy to specify specific flags through a build
system (like the necessary -Xclang -fembed-bitcode=all), but the build
system produces no manifest of compilation commands like CMake's
compile_commands.json. This isn't reccomended to use in place of a
proper manifest but is necessary for supporting build systems like Make.

Author: boomanaiden154

compiler_opt/tools/extract_ir.py

@@ -41,10 +41,12 @@
 
 flags.DEFINE_string(
     'input', None,
-    'Input file - either compile_commands.json or a linker parameter list')
+    'Input file or directory - either compile_commands.json, a linker parameter'
+    'list, or a path to a directory containing object files.')
 flags.DEFINE_enum(
-    'input_type', 'json', ['json', 'params'],
-    'Input file type - json or params. The latter refers to lld params.')
+    'input_type', 'json', ['json', 'params', 'directory'],
+    'Input file type - json, params, or directory. params latter refers to lld'
+    'params.')
 flags.DEFINE_string('output_dir', None, 'Output directory')
 flags.DEFINE_integer(
     'num_workers', None,
@@ -110,6 +112,13 @@ def main(argv):
       objs = extract_ir_lib.load_from_lld_params(
           [l.strip() for l in f.readlines()], FLAGS.obj_base_dir,
           FLAGS.output_dir)
+  elif FLAGS.input_type == 'directory':
+    logging.warning(
+        'Using the directory input is only reccomended if the build system'
+        'your project uses does not support any structured output that'
+        'ml-compiler-opt understands. If your build system provides a'
+        'structured compilation database, use that instead')
+    objs = extract_ir_lib.load_from_directory(FLAGS.input, FLAGS.output_dir)
   else:
     logging.error('Unknown input type: %s', FLAGS.input_type)
 

compiler_opt/tools/extract_ir_lib.py

@@ -250,6 +250,19 @@ def make_obj(obj_file: str) -> TrainingIRExtractor:
   return [make_obj(obj_file) for obj_file in just_obj_paths]
 
 
+def load_from_directory(obj_base_dir: str,
+                        output_dir: str) -> List[TrainingIRExtractor]:
+  paths = [str(p) for p in pathlib.Path(obj_base_dir).glob('**/*.o')]
+
+  def make_spec(obj_file: str):
+    return TrainingIRExtractor(
+        obj_relative_path=os.path.relpath(obj_file, start=obj_base_dir),
+        output_base_dir=output_dir,
+        obj_base_dir=obj_base_dir)
+
+  return [make_spec(path) for path in 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')

compiler_opt/tools/extract_ir_test.py

@@ -120,6 +120,21 @@ 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_load_from_directory(self):
+    tempdir = self.create_tempdir()
+    subdir = tempdir.mkdir(dir_path='subdir')
+    subdir.create_file(file_path='test1.o')
+    subdir.create_file(file_path='test2.o')
+    outdir = self.create_tempdir()
+    objs = extract_ir_lib.load_from_directory(tempdir.full_path,
+                                              outdir.full_path)
+    self.assertLen(objs, 2)
+    for index, obj in enumerate(
+        sorted(objs, key=lambda x: x._obj_relative_path)):
+      self.assertEqual(obj._obj_relative_path, f'subdir/test{index + 1:d}.o')
+      self.assertEqual(obj._obj_base_dir, tempdir.full_path)
+      self.assertEqual(obj._output_base_dir, outdir.full_path)
+
   def test_lld_thinlto_discovery(self):
     tempdir = self.create_tempdir()
     tempdir.create_file(file_path='1.3.import.bc')