1 : // Copyright 2013 Google Inc. All Rights Reserved.
2 : //
3 : // Licensed under the Apache License, Version 2.0 (the "License");
4 : // you may not use this file except in compliance with the License.
5 : // You may obtain a copy of the License at
6 : //
7 : // http://www.apache.org/licenses/LICENSE-2.0
8 : //
9 : // Unless required by applicable law or agreed to in writing, software
10 : // distributed under the License is distributed on an "AS IS" BASIS,
11 : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 : // See the License for the specific language governing permissions and
13 : // limitations under the License.
14 :
15 : #include "syzygy/optimize/optimize_app.h"
16 :
17 : #include "syzygy/block_graph/block_graph.h"
18 : #include "syzygy/block_graph/transforms/fuzzing_transform.h"
19 : #include "syzygy/block_graph/transforms/named_transform.h"
20 : #include "syzygy/common/indexed_frequency_data.h"
21 : #include "syzygy/grinder/basic_block_util.h"
22 : #include "syzygy/optimize/application_profile.h"
23 : #include "syzygy/optimize/transforms/basic_block_reordering_transform.h"
24 : #include "syzygy/optimize/transforms/block_alignment_transform.h"
25 : #include "syzygy/optimize/transforms/chained_subgraph_transforms.h"
26 : #include "syzygy/optimize/transforms/inlining_transform.h"
27 : #include "syzygy/optimize/transforms/peephole_transform.h"
28 : #include "syzygy/pe/pe_relinker.h"
29 : #include "syzygy/pe/pe_transform_policy.h"
30 :
31 : namespace optimize {
32 :
33 : namespace {
34 :
35 : using block_graph::transforms::FuzzingTransform;
36 : using common::IndexedFrequencyData;
37 : using grinder::basic_block_util::IndexedFrequencyMap;
38 : using grinder::basic_block_util::LoadBranchStatisticsFromFile;
39 : using optimize::transforms::BasicBlockReorderingTransform;
40 : using optimize::transforms::BlockAlignmentTransform;
41 : using optimize::transforms::ChainedSubgraphTransforms;
42 : using optimize::transforms::InliningTransform;
43 : using optimize::transforms::PeepholeTransform;
44 :
45 : const char kUsageFormatStr[] =
46 : "Usage: %ls [options]\n"
47 : " Required Options:\n"
48 : " --input-image=<path> The input image file to optimize.\n"
49 : " --output-image=<path> Output path for the rewritten image file.\n"
50 : "\n"
51 : " Options:\n"
52 : " --branch-file=<path> Branch statistics in JSON format.\n"
53 : " --input-pdb=<path> The PDB file associated with the input DLL.\n"
54 : " Default is inferred from input-image.\n"
55 : " --output-pdb=<path> Output path for the rewritten PDB file.\n"
56 : " Default is inferred from output-image.\n"
57 : " --overwrite Allow output files to be overwritten.\n"
58 : "\n"
59 : " Optimization Options:\n"
60 : " --all Enable all optimizations.\n"
61 : " --basic-block-reorder Enable basic block reodering.\n"
62 : " --block-alignment Enable block realignment.\n"
63 : " --inlining Enable function inlining.\n"
64 : " --peephole Enable peephole optimization.\n"
65 : "\n"
66 : " Testing Options:\n"
67 : " --fuzz Fuzz the binary.\n"
68 : "\n";
69 :
70 : } // namespace
71 :
72 E : bool OptimizeApp::ParseCommandLine(const CommandLine* cmd_line) {
73 :
74 E : if (cmd_line->HasSwitch("help"))
75 E : return Usage(cmd_line, "");
76 :
77 E : input_image_path_ = AbsolutePath(cmd_line->GetSwitchValuePath("input-image"));
78 E : output_image_path_ = cmd_line->GetSwitchValuePath("output-image");
79 E : input_pdb_path_ = AbsolutePath(cmd_line->GetSwitchValuePath("input-pdb"));
80 E : output_pdb_path_ = cmd_line->GetSwitchValuePath("output-pdb");
81 E : branch_file_path_ = AbsolutePath(cmd_line->GetSwitchValuePath("branch-file"));
82 :
83 E : basic_block_reorder_ = cmd_line->HasSwitch("basic-block-reorder");
84 E : block_alignment_ = cmd_line->HasSwitch("block-alignment");
85 E : fuzz_ = cmd_line->HasSwitch("fuzz");
86 E : inlining_ = cmd_line->HasSwitch("inlining");
87 E : peephole_ = cmd_line->HasSwitch("peephole");
88 E : overwrite_ = cmd_line->HasSwitch("overwrite");
89 :
90 : // Enable all optimization transforms.
91 E : if (cmd_line->HasSwitch("all")) {
92 E : basic_block_reorder_ = true;
93 E : block_alignment_ = true;
94 E : inlining_ = true;
95 E : peephole_ = true;
96 : }
97 :
98 : // The --input-image argument is required.
99 E : if (input_image_path_.empty())
100 E : return Usage(cmd_line, "You must specify --input-image.");
101 :
102 : // The --output-image argument is required.
103 E : if (output_image_path_.empty())
104 E : return Usage(cmd_line, "You must specify --output-image.");
105 :
106 E : return true;
107 E : }
108 :
109 E : bool OptimizeApp::SetUp() {
110 E : DCHECK(!input_image_path_.empty());
111 E : DCHECK(!output_image_path_.empty());
112 E : return true;
113 E : }
114 :
115 E : int OptimizeApp::Run() {
116 E : pe::PETransformPolicy policy;
117 E : pe::PERelinker relinker(&policy);
118 E : relinker.set_input_path(input_image_path_);
119 E : relinker.set_input_pdb_path(input_pdb_path_);
120 E : relinker.set_output_path(output_image_path_);
121 E : relinker.set_output_pdb_path(output_pdb_path_);
122 E : relinker.set_allow_overwrite(overwrite_);
123 :
124 : // Initialize the relinker. This does the decomposition, etc.
125 E : if (!relinker.Init()) {
126 i : LOG(ERROR) << "Failed to initialize relinker.";
127 i : return 1;
128 : }
129 :
130 : // Get module signature and layout.
131 E : pe::PEFile::Signature signature;
132 E : relinker.input_pe_file().GetSignature(&signature);
133 E : const pe::ImageLayout& image_layout = relinker.input_image_layout();
134 :
135 : // Load profile information from file.
136 E : ApplicationProfile profile(&image_layout);
137 E : if (!branch_file_path_.empty()) {
138 i : IndexedFrequencyMap frequencies;
139 : if (!LoadBranchStatisticsFromFile(branch_file_path_,
140 : signature,
141 i : &frequencies)) {
142 i : LOG(ERROR) << "Unable to load profile information.";
143 i : return 1;
144 : }
145 i : if (!profile.ImportFrequencies(frequencies)) {
146 i : LOG(ERROR) << "Could not import metrics for '"
147 : << branch_file_path_.value() << "'.";
148 i : return false;
149 : }
150 i : }
151 :
152 : // Compute global profile information for the current block graph.
153 E : if (!profile.ComputeGlobalProfile()) {
154 i : LOG(ERROR) << "Unable to build profile information.";
155 i : return 1;
156 : }
157 :
158 : // Construct a chain of basic block transforms.
159 E : ChainedSubgraphTransforms chains(&profile);
160 :
161 : // Declare transforms we may apply.
162 E : scoped_ptr<BasicBlockReorderingTransform> basic_block_reordering_transform;
163 E : scoped_ptr<BlockAlignmentTransform> block_alignment_transform;
164 E : scoped_ptr<FuzzingTransform> fuzzing_transform;
165 E : scoped_ptr<InliningTransform> inlining_transform;
166 E : scoped_ptr<PeepholeTransform> peephole_transform;
167 :
168 : // If block block reordering is enabled, add it to the chain.
169 E : if (peephole_) {
170 i : peephole_transform.reset(new PeepholeTransform());
171 i : chains.AppendTransform(peephole_transform.get());
172 : }
173 :
174 : // If inlining is enabled, add it to the chain.
175 E : if (inlining_) {
176 i : inlining_transform.reset(new InliningTransform());
177 i : chains.AppendTransform(inlining_transform.get());
178 : }
179 :
180 : // If block block reordering is enabled, add it to the chain.
181 E : if (basic_block_reorder_) {
182 i : basic_block_reordering_transform.reset(new BasicBlockReorderingTransform());
183 i : chains.AppendTransform(basic_block_reordering_transform.get());
184 : }
185 :
186 : // If block alignment is enabled, add it to the chain.
187 E : if (block_alignment_) {
188 i : block_alignment_transform.reset(new BlockAlignmentTransform());
189 i : chains.AppendTransform(block_alignment_transform.get());
190 : }
191 :
192 : // Append the chain to the relinker.
193 E : if (!relinker.AppendTransform(&chains))
194 i : return false;
195 :
196 : // If fuzzing is enabled, add it to the relinker.
197 E : if (fuzz_) {
198 i : fuzzing_transform.reset(new block_graph::transforms::FuzzingTransform);
199 i : relinker.AppendTransform(fuzzing_transform.get());
200 : }
201 :
202 : // Perform the actual relink.
203 E : if (!relinker.Relink()) {
204 i : LOG(ERROR) << "Unable to relink input image.";
205 i : return 1;
206 : }
207 :
208 E : return 0;
209 E : }
210 :
211 : bool OptimizeApp::Usage(const CommandLine* cmd_line,
212 E : const base::StringPiece& message) const {
213 E : if (!message.empty()) {
214 E : ::fwrite(message.data(), 1, message.length(), err());
215 E : ::fprintf(err(), "\n\n");
216 : }
217 :
218 : ::fprintf(err(),
219 : kUsageFormatStr,
220 E : cmd_line->GetProgram().BaseName().value().c_str());
221 :
222 E : return false;
223 E : }
224 :
225 : } // namespace optimize
|