Add make_corpus utility

This patch adds in a make_corpus.py script (along with the associated
library and unit test) that allows for the creation of a corpus from a
directory containing LLVM bitcode files.

Author: boomanaiden154

compiler_opt/tools/make_corpus.py

@@ -0,0 +1,57 @@
+# coding=utf-8
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tool for making a corpus from arbitrary bitcode.
+
+To create a corpus from a set of bitcode files in an input directory, run
+the following command:
+
+PYTHONPATH=$PYTHONPATH:. python3 ./compiler_opt/tools/make_corpus.py \
+  --input_dir=<path to input directory> \
+  --output_dir=<path to output direcotry> \
+  --default_flags="<list of space separated flags>"
+"""
+
+from absl import app
+from absl import flags
+from absl import logging
+
+from compiler_opt.tools import make_corpus_lib
+
+flags.DEFINE_string('input_dir', None, 'The input directory.')
+flags.DEFINE_string('output_dir', None, 'The output directory.')
+flags.DEFINE_string(
+    'default_flags', '',
+    'The compiler flags to compile with when using downstream tooling.')
+
+flags.mark_flag_as_required('input_dir')
+flags.mark_flag_as_required('output_dir')
+
+FLAGS = flags.FLAGS
+
+
+def main(_):
+  logging.warn('Using this tool does not guarnatee that the bitcode is taken at'
+               'the correct stage for consumption during model training. Make'
+               'sure to validate assumptions about where the bitcode is coming'
+               'from before using it in production.')
+  relative_paths = make_corpus_lib.load_bitcode_from_directory(FLAGS.input_dir)
+  make_corpus_lib.copy_bitcode(relative_paths, FLAGS.input_dir,
+                               FLAGS.output_dir)
+  make_corpus_lib.write_corpus_manifest(relative_paths, FLAGS.output_dir,
+                                        FLAGS.default_flags.split())
+
+
+if __name__ == '__main__':
+  app.run(main)

compiler_opt/tools/make_corpus_lib.py

@@ -0,0 +1,83 @@
+# coding=utf-8
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Library functions for making a corpus from arbitrary bitcode."""
+
+import pathlib
+import os
+import shutil
+import json
+
+from typing import List, Optional
+
+
+def load_bitcode_from_directory(bitcode_base_dir: str):
+  """Finds bitcode files to extract from a given directory.
+
+  Args:
+    bitcode_base_dir: The base directory where the bitcode to be copied
+      is from.
+    output_dir: The directory to place the bitcode in.
+
+  Returns an array of paths representing the relative path to the bitcode
+  file from the base direcotry.
+  """
+  paths = [str(p) for p in pathlib.Path(bitcode_base_dir).glob('**/*.bc')]
+
+  return [
+      os.path.relpath(full_path, start=bitcode_base_dir) for full_path in paths
+  ]
+
+
+def copy_bitcode(relative_paths: List[str], bitcode_base_dir: str,
+                 output_dir: str):
+  """Copies bitcode files from the base directory to the output directory.
+
+  Args:
+    relative_paths: An array of relative paths to bitcode files that are copied
+      over to the output directory, preserving relative location.
+    bitcode_base_dir: The base directory where the bitcode is located.
+    output_dir: The output directory to place the bitcode in.
+  """
+  for relative_path in relative_paths:
+    base_path = os.path.join(bitcode_base_dir, relative_path)
+    destination_path = os.path.join(output_dir, relative_path)
+    os.makedirs(os.path.dirname(destination_path), exist_ok=True)
+    shutil.copy(base_path, destination_path)
+
+
+def write_corpus_manifest(relative_output_paths: List[str],
+                          output_dir: str,
+                          default_flags: Optional[List[str]] = None):
+  """Creates a corpus manifest describing the bitcode that has been found.
+
+  Args:
+    relative_output_paths: A list of paths to each bitcode file relative to the
+      output directory.
+    outout_dir: The output directory where the corpus is being created.
+    default_flags: An array of compiler flags that should be used to compile
+      the bitcode when using further downstream tooling."""
+  if default_flags is None:
+    default_flags = []
+  corpus_description = {
+      'global_command_override': default_flags,
+      'has_thinlto': False,
+      'modules': [path for path in relative_output_paths if path is not None]
+  }
+
+  with open(
+      os.path.join(output_dir, 'corpus_description.json'),
+      'w',
+      encoding='utf-8') as description_file:
+    json.dump(corpus_description, description_file, indent=2)

compiler_opt/tools/make_corpus_test.py

@@ -0,0 +1,66 @@
+# coding=utf-8
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Test for compiler_opt.tools.make_corpus_lib"""
+
+import json
+import os
+
+from absl.testing import absltest
+
+from compiler_opt.tools import make_corpus_lib
+
+
+class MakeCorpusTest(absltest.TestCase):
+
+  def test_load_bitcode_from_directory(self):
+    outer = self.create_tempdir()
+    tempdir = outer.mkdir(dir_path='nested')
+    tempdir.create_file('test1.bc')
+    tempdir.create_file('test2.bc')
+    relative_paths = make_corpus_lib.load_bitcode_from_directory(outer)
+    relative_paths = sorted(relative_paths)
+    self.assertEqual(relative_paths[0], 'nested/test1.bc')
+    self.assertEqual(relative_paths[1], 'nested/test2.bc')
+
+  def test_copy_bitcode(self):
+    build_dir = self.create_tempdir()
+    nested_dir = build_dir.mkdir(dir_path='nested')
+    nested_dir.create_file('test1.bc')
+    nested_dir.create_file('test2.bc')
+    relative_paths = ['nested/test1.bc', 'nested/test2.bc']
+    corpus_dir = self.create_tempdir()
+    make_corpus_lib.copy_bitcode(relative_paths, build_dir, corpus_dir)
+    output_files = sorted(os.listdir(os.path.join(corpus_dir, './nested')))
+    self.assertEqual(output_files[0], 'test1.bc')
+    self.assertEqual(output_files[1], 'test2.bc')
+
+  def test_write_corpus_manifest(self):
+    relative_output_paths = ['test/test1.bc', 'test/test2.bc']
+    output_dir = self.create_tempdir()
+    default_flags = ['-O3', '-c']
+    make_corpus_lib.write_corpus_manifest(relative_output_paths, output_dir,
+                                          default_flags)
+    with open(
+        os.path.join(output_dir, 'corpus_description.json'),
+        encoding='utf-8') as corpus_description_file:
+      corpus_description = json.load(corpus_description_file)
+    self.assertEqual(corpus_description['global_command_override'],
+                     default_flags)
+    self.assertEqual(corpus_description['has_thinlto'], False)
+    self.assertEqual(corpus_description['modules'], relative_output_paths)
+
+
+if __name__ == '__main__':
+  absltest.main()