# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Extract IR for training.
Extract IR for training, either from a compile_commands.json file produced by
cmake, or a linker parameter list file.
Only run with
'python compiler_opt/tools/extract_ir.py ...'
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 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
To change the logging verbosity, set the --verbosity flag to the desired level.
Setting it to a specific level will enable all messages at that level and
higher. Exact values can be found by invoking the script with --help.
"""
import argparse
import json
import logging
from mlgo.corpus import extract_ir_lib
def parse_args_and_run():
parser = argparse.ArgumentParser(
description="A tool for making a corpus from build artifacts"
)
parser.add_argument(
"--input",
type=str,
help="Input file or directory - either compile_commands.json, a linker "
"parameter list, or a path to a directory containing object files.",
)
parser.add_argument(
"--input_type",
type=str,
help="Input file type - JSON, LLD params, directory, or bazel aquery.",
choices=["json", "params", "directory", "bazel_aquery"],
default="json",
nargs="?",
)
parser.add_argument("--output_dir", type=str, help="Output directory")
parser.add_argument(
"--num_workers",
type=int,
help="Number of parallel works for objcopy. `None` for maximum available.",
default=None,
nargs="?",
)
parser.add_argument(
"--llvm_objcopy_path",
type=str,
help="Path to llvm-objcopy",
default="llvm-objcopy",
nargs="?",
)
parser.add_argument(
"--obj_base_dir",
type=str,
help="Base directory for object files. Defaults to current working dir.",
default="",
nargs="?",
)
parser.add_argument(
"--cmd_filter",
type=str,
help="Include only those modules with a command line matching this regular "
"expression. Set it to None to not perform any filtering. Note that the "
"regular expression is applied independently for each separate command line "
"option. For example, ^-Oz$ will match Oz built binaries. This does not work "
"with thinlto_build=lld.",
default=None,
nargs="?",
)
parser.add_argument(
"--thinlto_build",
type=str,
help="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",
choices=["distributed", "local"],
default=None,
nargs="?",
)
parser.add_argument(
"--cmd_section_name",
type=str,
help="The section name passed to llvm-objcopy. For ELF object files, the "
"default .llvmcmd is correct. For Mach-O object files, one should use "
"something like __LLVM,__cmdline",
default=".llvmcmd",
nargs="?",
)
parser.add_argument(
"--bitcode_section_name",
type=str,
help="The section name passed to llvm-objcopy. For ELF object files, the "
"default .llvmbc is correct. For Mach-O object files, one should use "
"__LLVM,__bitcode",
default=".llvmbc",
nargs="?",
)
# TODO(#107898): Refactor this into a common location.
parser.add_argument(
"--verbosity",
type=str,
help="The verbosity level to use for logging",
default="INFO",
nargs="?",
choices=["DEBUG", "INFO", "WARNING", "ERROR"],
)
args = parser.parse_args()
main(args)
def main(args):
logging.basicConfig(level=args.verbosity)
objs = []
if args.input is not None and args.thinlto_build == "local":
raise ValueError("--thinlto_build=local cannot be run with --input")
if args.input is None:
if args.thinlto_build != "local":
raise ValueError("--input or --thinlto_build=local must be provided")
objs = extract_ir_lib.load_for_lld_thinlto(args.obj_base_dir, args.output_dir)
elif args.input_type == "json":
with open(args.input, encoding="utf-8") as f:
objs = extract_ir_lib.load_from_compile_commands(
json.load(f), args.output_dir
)
elif args.input_type == "params":
if not args.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(args.input, encoding="utf-8") as f:
objs = extract_ir_lib.load_from_lld_params(
[l.strip() for l in f.readlines()], args.obj_base_dir, args.output_dir
)
elif args.input_type == "directory":
logging.warning(
"Using the directory input is only recommended 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(args.input, args.output_dir)
elif args.input_type == "bazel_aquery":
with open(args.input, encoding="utf-8") as aquery_json_handle:
objs = extract_ir_lib.load_bazel_aquery(
json.load(aquery_json_handle), args.obj_base_dir, args.output_dir
)
else:
logging.error("Unknown input type: %s", args.input_type)
relative_output_paths = extract_ir_lib.run_extraction(
objs,
args.num_workers,
args.llvm_objcopy_path,
args.cmd_filter,
args.thinlto_build,
args.cmd_section_name,
args.bitcode_section_name,
)
extract_ir_lib.write_corpus_manifest(
args.thinlto_build, relative_output_paths, args.output_dir
)
logging.info(
"Converted %d files out of %d",
len(objs) - relative_output_paths.count(None),
len(objs),
)
if __name__ == "__main__":
parse_args_and_run()
Codebase Browser
llvm