trace_gen.rs 23 KB
Newer Older
Isaac Oscar Gariano's avatar
Isaac Oscar Gariano committed
1
// Copyright 2017 The Australian National University
qinsoon's avatar
qinsoon committed
2
//
Isaac Oscar Gariano's avatar
Isaac Oscar Gariano committed
3
4
5
// 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
qinsoon's avatar
qinsoon committed
6
//
Isaac Oscar Gariano's avatar
Isaac Oscar Gariano committed
7
//     http://www.apache.org/licenses/LICENSE-2.0
qinsoon's avatar
qinsoon committed
8
//
Isaac Oscar Gariano's avatar
Isaac Oscar Gariano committed
9
10
11
12
13
14
// 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.

15
use ast::ir::*;
16
17
18
19
use ast::types::*;
use ast::inst::*;
use ast::ptr::*;
use ast::op::CmpOp;
qinsoon's avatar
qinsoon committed
20
use vm::VM;
21
use compiler::CompilerPass;
22
use utils::LinkedHashSet;
qinsoon's avatar
qinsoon committed
23
24
use std::any::Any;

25
pub struct TraceGen {
26
    name: &'static str
27
28
29
30
}

impl TraceGen {
    pub fn new() -> TraceGen {
qinsoon's avatar
qinsoon committed
31
        TraceGen {
32
            name: "Trace Generation"
qinsoon's avatar
qinsoon committed
33
        }
34
35
36
    }
}

qinsoon's avatar
qinsoon committed
37
const LOG_TRACE_SCHEDULE: bool = true;
38

39
40
41
42
impl CompilerPass for TraceGen {
    fn name(&self) -> &'static str {
        self.name
    }
qinsoon's avatar
qinsoon committed
43
44
45
46

    fn as_any(&self) -> &Any {
        self
    }
qinsoon's avatar
qinsoon committed
47

48
    #[allow(unused_variables)] // vm is not used here
qinsoon's avatar
qinsoon committed
49
    fn visit_function(&mut self, vm: &VM, func: &mut MuFunctionVersion) {
50
51
52
        // we put the high probability edge into a hot trace, and others into cold paths
        // and traverse cold_path later
        let trace = {
qinsoon's avatar
qinsoon committed
53
            let mut trace: Vec<MuID> = vec![];
54
55

            // main work stack
qinsoon's avatar
qinsoon committed
56
            let mut work_stack: LinkedHashSet<MuID> = LinkedHashSet::new();
57
            // slow path queue (they are scheduled after main work stack is finished)
qinsoon's avatar
qinsoon committed
58
            let mut slowpath_queue: LinkedHashSet<MuID> = LinkedHashSet::new();
59
            // return sink (always schedule this after all blocks)
qinsoon's avatar
qinsoon committed
60
            let mut ret_sink: Option<MuID> = None;
61
62
63
64

            let f_content = func.content.as_ref().unwrap();
            let entry = f_content.entry;
            work_stack.insert(entry);
qinsoon's avatar
qinsoon committed
65

66
            while !work_stack.is_empty() || !slowpath_queue.is_empty() {
qinsoon's avatar
qinsoon committed
67
                let cur_block: &Block = {
68
69
70
71
72
73
74
75
                    let ret = if let Some(b) = work_stack.pop_back() {
                        b
                    } else if let Some(b) = slowpath_queue.pop_front() {
                        b
                    } else {
                        unreachable!()
                    };
                    f_content.get_block(ret)
76
                };
qinsoon's avatar
qinsoon committed
77
78
79
80
81
82
                trace_if!(
                    LOG_TRACE_SCHEDULE,
                    "---check block {} #{}---",
                    cur_block,
                    cur_block.id()
                );
83
84

                // append current block to the trace
85
                if !trace.contains(&cur_block.id()) {
qinsoon's avatar
qinsoon committed
86
87
88
89
90
91
                    trace_if!(
                        LOG_TRACE_SCHEDULE,
                        "add {} #{} to trace",
                        cur_block,
                        cur_block.id()
                    );
92
                    trace.push(cur_block.id());
93

94
95
96
                    // trying to find next block
                    let next_block: Option<MuID> = match find_next_block(cur_block, func) {
                        Some(id) => Some(id),
97
                        None => None
98
                    };
qinsoon's avatar
qinsoon committed
99
100
101
102
103
104
                    trace_if!(
                        LOG_TRACE_SCHEDULE && next_block.is_some(),
                        "find next block as {} #{}",
                        f_content.get_block(next_block.unwrap()),
                        next_block.unwrap()
                    );
105
106

                    // put other succeeding blocks to different work stacks
qinsoon's avatar
qinsoon committed
107
108
109
110
111
112
                    let mut all_successors: LinkedHashSet<MuID> = LinkedHashSet::from_vec(
                        cur_block
                            .control_flow
                            .succs
                            .iter()
                            .map(|x| x.target)
113
                            .collect()
qinsoon's avatar
qinsoon committed
114
                    );
115
116
117
118
119
120
121
122
123
124
                    // remove next block from it
                    if next_block.is_some() {
                        all_successors.remove(&next_block.unwrap());
                    }

                    // push other successors to different work queues
                    for succ_id in all_successors.iter() {
                        let succ = f_content.get_block(*succ_id);
                        match succ.trace_hint {
                            TraceHint::None => {
qinsoon's avatar
qinsoon committed
125
126
127
128
129
130
                                trace_if!(
                                    LOG_TRACE_SCHEDULE,
                                    "push {} #{} to work stack",
                                    succ,
                                    succ_id
                                );
131
132
133
                                work_stack.insert(*succ_id);
                            }
                            TraceHint::SlowPath => {
qinsoon's avatar
qinsoon committed
134
135
136
137
138
139
                                trace_if!(
                                    LOG_TRACE_SCHEDULE,
                                    "push {} #{} to slow path",
                                    succ,
                                    succ_id
                                );
140
141
142
                                slowpath_queue.insert(*succ_id);
                            }
                            TraceHint::ReturnSink => {
qinsoon's avatar
qinsoon committed
143
144
145
146
147
148
149
150
151
152
153
                                assert!(
                                    ret_sink.is_none() ||
                                        (ret_sink.is_some() && ret_sink.unwrap() == *succ_id),
                                    "cannot have more than one return sink"
                                );
                                trace_if!(
                                    LOG_TRACE_SCHEDULE,
                                    "set {} #{} as return sink",
                                    succ,
                                    succ_id
                                );
154
155
156
                                ret_sink = Some(*succ_id);
                            }
                            TraceHint::FastPath => {
qinsoon's avatar
qinsoon committed
157
158
159
160
161
162
163
164
165
                                panic!(
                                    "trying to delay the insertion of a block with fastpath hint: \
                                     {} #{}. Either we missed to pick it as next block, or the \
                                     current checking block has several succeeding blocks with \
                                     fastpath hint which is \
                                     not reasonable",
                                    succ,
                                    succ_id
                                );
166
                            }
167
168
169
                        }
                    }

170
171
172
173
                    // push next block to work stack - this comes after the pushes above, so it gets
                    // popped earlier (and scheduled before those)
                    if next_block.is_some() {
                        let next_block = next_block.unwrap();
qinsoon's avatar
qinsoon committed
174
175
176
177
178
179
                        trace_if!(
                            LOG_TRACE_SCHEDULE,
                            "push hot edge {} #{} to work stack",
                            f_content.get_block(next_block),
                            next_block
                        );
180
181
182
183
                        work_stack.insert(next_block);
                    }

                    trace_if!(LOG_TRACE_SCHEDULE, "");
184
                } else {
185
186
                    trace_if!(LOG_TRACE_SCHEDULE, "block already in trace, ignore");
                    continue;
187
188
                }
            }
189
190
191

            // add return sink
            if let Some(ret_sink) = ret_sink {
qinsoon's avatar
qinsoon committed
192
193
194
195
196
197
198
199
200
201
                assert!(
                    !trace.contains(&ret_sink),
                    "return sink should not already be scheduled"
                );
                trace_if!(
                    LOG_TRACE_SCHEDULE,
                    "push return sink {} #{} to the trace",
                    f_content.get_block(ret_sink),
                    ret_sink
                );
202
203
204
                trace.push(ret_sink);
            }

205
206
            trace
        };
qinsoon's avatar
qinsoon committed
207

208
209
        func.block_trace = Some(trace);
    }
qinsoon's avatar
qinsoon committed
210

211
    #[allow(unused_variables)] // vm is not used here
qinsoon's avatar
qinsoon committed
212
    fn finish_function(&mut self, vm: &VM, func: &mut MuFunctionVersion) {
qinsoon's avatar
qinsoon committed
213
        debug!("trace for {}", func);
214
        debug!("{:?}", func.block_trace.as_ref().unwrap());
215
216
217
218
219
220

        info!("doing branch adjustment...");
        branch_adjustment(func, vm);

        debug!("trace for {}", func);
        debug!("{:?}", func.block_trace.as_ref().unwrap());
221
    }
222
}
223
224
225
226
227
228
229
230
231

/// returns the successor of current block.
/// We first look at trace hint, if there is no trace hint to indicate next block,
/// we layout the block with the highest probability as next block (in case of a tie,
/// returns the first met successor). If current block does not have any successor,
/// returns None.
fn find_next_block(cur_block: &Block, func: &MuFunctionVersion) -> Option<MuID> {
    let f_content = func.content.as_ref().unwrap();
    let ref succs = cur_block.control_flow.succs;
qinsoon's avatar
qinsoon committed
232
233
234
    let has_fastpath = succs.iter().find(|edge| {
        f_content.get_block(edge.target).trace_hint == TraceHint::FastPath
    });
235
236

    if has_fastpath.is_some() {
qinsoon's avatar
qinsoon committed
237
        let target = has_fastpath.unwrap().target;
qinsoon's avatar
qinsoon committed
238
239
240
241
242
243
        trace_if!(
            LOG_TRACE_SCHEDULE,
            "found fastpath successor {} for block {}",
            target,
            cur_block
        );
qinsoon's avatar
qinsoon committed
244
        Some(target)
245
246
247
    } else {
        // we need to find next path by examining probability
        if succs.len() == 0 {
qinsoon's avatar
qinsoon committed
248
249
250
251
252
            trace_if!(
                LOG_TRACE_SCHEDULE,
                "cannot find successors of block {}",
                cur_block
            );
253
254
            None
        } else {
qinsoon's avatar
qinsoon committed
255
            trace_if!(LOG_TRACE_SCHEDULE, "successors: {:?}", succs);
qinsoon's avatar
qinsoon committed
256
257
258
            let ideal_successors: Vec<&BlockEdge> = succs
                .iter()
                .filter(|b| match f_content.get_block(b.target).trace_hint {
259
                    TraceHint::SlowPath | TraceHint::ReturnSink => false,
260
                    _ => true
qinsoon's avatar
qinsoon committed
261
262
263
264
265
266
267
                })
                .collect();
            trace_if!(
                LOG_TRACE_SCHEDULE,
                "after filtering out slowpath/retsink, we have: {:?}",
                ideal_successors
            );
268
269
270
271

            if ideal_successors.len() == 0 {
                None
            } else {
qinsoon's avatar
qinsoon committed
272
273
                let mut hot_blk = ideal_successors[0].target;
                let mut hot_prob = ideal_successors[0].probability;
274

qinsoon's avatar
qinsoon committed
275
                for edge in ideal_successors.iter() {
qinsoon's avatar
qinsoon committed
276
277
278
279
280
281
                    trace_if!(
                        LOG_TRACE_SCHEDULE,
                        "succ: {}/{}",
                        edge.target,
                        edge.probability
                    );
qinsoon's avatar
qinsoon committed
282
                    if edge.probability >= hot_prob {
283
284
285
286
287
288
289
290
291
292
293
                        hot_blk = edge.target;
                        hot_prob = edge.probability;
                    }
                }

                Some(hot_blk)
            }
        }
    }
}

qinsoon's avatar
qinsoon committed
294
295
/// a conditional branch should always be followed by its false label.
/// The adjustment should follow the rules:
296
297
298
299
300
301
302
303
304
/// * any conditional branch followed by its false label stays unchanged
/// * for conditional branch followed by its true label,
///   we switch the true and false label, and negate the condition
/// * for conditional branch followed by neither label,
///   we invent a new false label, and rewrite the conditional branch so that
///   the new cond branch will be followed by the new false label.
fn branch_adjustment(func: &mut MuFunctionVersion, vm: &VM) {
    let mut trace = func.block_trace.take().unwrap();
    let mut f_content = func.content.take().unwrap();
qinsoon's avatar
qinsoon committed
305
    let mut new_blocks: Vec<Block> = vec![];
306

307
308
309
    for (blk_id, mut block) in f_content.blocks.iter_mut() {
        trace_if!(LOG_TRACE_SCHEDULE, "block: {} #{}", block, blk_id);

qinsoon's avatar
qinsoon committed
310
        let next_block_in_trace: Option<usize> = {
311
312
313
314
315
            if let Some(index) = trace.iter().position(|x| x == blk_id) {
                if index == trace.len() {
                    None
                } else {
                    Some(index + 1)
316
                }
317
318
            } else {
                warn!("find an unreachable block (a block exists in IR, but is not in trace");
qinsoon's avatar
qinsoon committed
319
                continue;
320
            }
321
322
323
324
325
326
327
328
329
330
331
332
        };

        // we only need to deal with blocks that ends with a branch
        if block.ends_with_cond_branch() {
            // old block content
            let block_content = block.content.as_ref().unwrap().clone();

            let mut new_body = vec![];

            for node in block_content.body.iter() {
                match node.v {
                    TreeNode_::Instruction(Instruction {
qinsoon's avatar
qinsoon committed
333
334
335
336
337
                        ref ops,
                        v: Instruction_::Branch2 {
                            cond,
                            ref true_dest,
                            ref false_dest,
338
                            true_prob
qinsoon's avatar
qinsoon committed
339
340
341
                        },
                        ..
                    }) => {
342
343
                        trace_if!(LOG_TRACE_SCHEDULE, "rewrite cond branch: {}", node);

qinsoon's avatar
qinsoon committed
344
                        let true_label = true_dest.target;
345
346
                        let false_label = false_dest.target;

qinsoon's avatar
qinsoon committed
347
348
349
                        if next_block_in_trace.is_some() &&
                            next_block_in_trace.unwrap() == false_label
                        {
350
351
352
                            // any conditional branch followed by its false label stays unchanged
                            trace_if!(LOG_TRACE_SCHEDULE, ">>stays unchanged");
                            new_body.push(node.clone());
qinsoon's avatar
qinsoon committed
353
354
355
                        } else if next_block_in_trace.is_some() &&
                                   next_block_in_trace.unwrap() == true_label
                        {
356
357
                            // for conditional branch followed by its true label
                            // we switch the true and false label, and negate the condition
qinsoon's avatar
qinsoon committed
358
                            let new_true_dest = false_dest.clone();
359
360
361
362
363
364
365
366
367
368
                            let new_false_dest = true_dest.clone();

                            let new_cond_node = {
                                let old_cond_node_clone = ops[cond].clone();
                                match ops[cond].v {
                                    // cond is a comparison, we recreate a comparison node
                                    // with inverted operator
                                    // orig: if a  OP b then L1 else L2
                                    // new : if a ~OP b then L2 else L1
                                    TreeNode_::Instruction(Instruction {
qinsoon's avatar
qinsoon committed
369
370
371
372
373
                                        ref value,
                                        ref ops,
                                        v: Instruction_::CmpOp(optr, op1, op2),
                                        ..
                                    }) => {
374
375
376
377
                                        TreeNode::new_inst(Instruction {
                                            hdr: MuEntityHeader::unnamed(vm.next_id()),
                                            value: value.clone(),
                                            ops: ops.clone(),
378
                                            v: Instruction_::CmpOp(optr.invert(), op1, op2)
379
380
381
382
383
384
385
                                        })
                                    }
                                    // cond is computed form other instruction or is a value
                                    // we add an instruction for cond EQ 0 (negate of cond EQ 1)
                                    // orig: if (cond)        then L1 else L2
                                    // new : if ((cond) EQ 0) then L2 else L1
                                    _ => {
386
                                        let temp_res: P<TreeNode> = func.new_ssa(
387
                                            MuEntityHeader::unnamed(vm.next_id()),
388
                                            UINT1_TYPE.clone()
389
390
391
392
                                        );
                                        let const_0 = func.new_constant(Value::make_int_const_ty(
                                            vm.next_id(),
                                            UINT1_TYPE.clone(),
393
                                            0
394
395
396
397
398
                                        ));
                                        TreeNode::new_inst(Instruction {
                                            hdr: MuEntityHeader::unnamed(vm.next_id()),
                                            value: Some(vec![temp_res.clone_value()]),
                                            ops: vec![old_cond_node_clone, const_0],
399
                                            v: Instruction_::CmpOp(CmpOp::EQ, 0, 1)
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
                                        })
                                    }
                                }
                            };

                            // old ops is something like: a b cond c d e
                            // we want: a b new_cond c d e
                            let new_ops = {
                                let mut ret = ops.clone();
                                ret.push(new_cond_node);
                                let last = ret.len() - 1;
                                ret.swap(cond, last);
                                ret.remove(last);
                                ret
                            };

                            let new_cond_branch = func.new_inst(Instruction {
                                hdr: MuEntityHeader::unnamed(vm.next_id()),
                                value: None,
                                ops: new_ops,
                                v: Instruction_::Branch2 {
                                    cond: cond,
                                    true_dest: new_true_dest,
                                    false_dest: new_false_dest,
424
425
                                    true_prob: 1f32 - true_prob
                                }
426
427
428
429
430
431
432
                            });

                            trace_if!(LOG_TRACE_SCHEDULE, ">>T/F labels switched, op negated");
                            trace_if!(LOG_TRACE_SCHEDULE, ">>{}", new_cond_branch);
                            new_body.push(new_cond_branch);
                        } else {
                            // for conditional branch followed by neither label,
qinsoon's avatar
qinsoon committed
433
434
                            // we invent a new false label, and rewrite the conditional branch
                            // so that the new cond branch will be followed by the new false label
435
436
437
438
439

                            // create a false block
                            // Lnew_false (arg list):
                            //   BRANCH Lfalse (arg list)
                            let new_false_block = {
440
441
                                let block_name =
                                    Arc::new(format!("{}:#{}:false", func.name(), node.id()));
qinsoon's avatar
qinsoon committed
442
443
                                let mut block =
                                    Block::new(MuEntityHeader::named(vm.next_id(), block_name));
444
445
                                vm.set_name(block.as_entity());

qinsoon's avatar
qinsoon committed
446
447
448
                                let block_args: Vec<P<TreeNode>> = false_dest
                                    .args
                                    .iter()
449
450
451
452
                                    .map(|x| match x {
                                        &DestArg::Normal(i) => {
                                            func.new_ssa(
                                                MuEntityHeader::unnamed(vm.next_id()),
453
                                                ops[i].as_value().ty.clone()
454
455
                                            )
                                        }
456
                                        _ => unimplemented!()
qinsoon's avatar
qinsoon committed
457
458
                                    })
                                    .collect();
459
460
461
462
463
464
465
466
467
468
469
470
                                let block_args_len = block_args.len();
                                block.content = Some(BlockContent {
                                    args: block_args.iter().map(|x| x.clone_value()).collect(),
                                    exn_arg: None,
                                    body: vec![
                                        func.new_inst(Instruction {
                                            hdr: MuEntityHeader::unnamed(vm.next_id()),
                                            value: None,
                                            ops: block_args,
                                            v: Instruction_::Branch1(Destination {
                                                target: false_dest.target,
                                                args: (0..block_args_len)
qinsoon's avatar
qinsoon committed
471
                                                    .map(|x| DestArg::Normal(x))
472
473
                                                    .collect()
                                            })
qinsoon's avatar
qinsoon committed
474
                                        }),
475
                                    ],
476
                                    keepalives: None
477
                                });
478

479
480
481
482
483
484
485
486
487
488
489
490
491
492
                                block
                            };

                            // for BRANCH2 cond Ltrue Lfalse
                            // rewrite it as BRANCH2 cond Ltrue (...) Lnew_false (...)
                            let new_cond_branch = func.new_inst(Instruction {
                                hdr: MuEntityHeader::unnamed(vm.next_id()),
                                value: None,
                                ops: ops.clone(),
                                v: Instruction_::Branch2 {
                                    cond: cond,
                                    true_dest: true_dest.clone(),
                                    false_dest: Destination {
                                        target: new_false_block.id(),
493
                                        args: false_dest.args.clone()
494
                                    },
495
496
                                    true_prob: true_prob
                                }
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
                            });

                            trace_if!(LOG_TRACE_SCHEDULE, ">>new F label created");
                            trace_if!(LOG_TRACE_SCHEDULE, ">>{}", new_cond_branch);
                            new_body.push(new_cond_branch);

                            // add new false block to trace (immediate after this block)
                            if let Some(next_block_index) = next_block_in_trace {
                                trace.insert(next_block_index, new_false_block.id());
                            } else {
                                trace.push(new_false_block.id());
                            }
                            // add new false block to new_blocks (insert them to function later)
                            new_blocks.push(new_false_block);
                        }
                    }

514
                    _ => new_body.push(node.clone())
515
516
517
518
519
520
521
                }
            }

            block.content = Some(BlockContent {
                args: block_content.args.to_vec(),
                exn_arg: block_content.exn_arg.clone(),
                body: new_body,
522
                keepalives: block_content.keepalives.clone()
523
            });
524
525
        }
    }
526
527
528
529
530
531
532

    for blk in new_blocks {
        f_content.blocks.insert(blk.id(), blk);
    }

    func.content = Some(f_content);
    func.block_trace = Some(trace);
qinsoon's avatar
qinsoon committed
533
}