1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
|
# 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, pass an integer representing the desired
verbosity to the --verbosity flag. Use 0 for all logs, status information,
and detailed debug information, -1 for solely warnings, and -2 to not produce
any output.
"""
import argparse
import json
import logging
import multiprocessing
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="?",
)
args = parser.parse_args()
main(args)
def main(args):
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()
|
