duskphantom_middle/transform/block_fuse.rs
1// Copyright 2024 Duskphantom Authors
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// SPDX-License-Identifier: Apache-2.0
16
17use anyhow::{Ok, Result};
18
19use crate::{
20 ir::{BBPtr, FunPtr},
21 Program,
22};
23
24use super::Transform;
25
26pub fn optimize_program(program: &mut Program) -> Result<bool> {
27 BlockFuse::new(program).run_and_log()
28}
29
30pub struct BlockFuse<'a> {
31 program: &'a mut Program,
32}
33
34impl<'a> Transform for BlockFuse<'a> {
35 fn get_program_mut(&mut self) -> &mut Program {
36 self.program
37 }
38
39 fn name() -> String {
40 "block_fuse".to_string()
41 }
42
43 fn run(&mut self) -> Result<bool> {
44 let mut changed = false;
45 for func in self.program.module.functions.clone() {
46 if func.is_lib() {
47 continue;
48 }
49 for bb in func.rpo_iter() {
50 changed |= self.fuse_block(bb, func)?;
51 }
52 }
53 Ok(changed)
54 }
55}
56
57impl<'a> BlockFuse<'a> {
58 pub fn new(program: &'a mut Program) -> Self {
59 Self { program }
60 }
61
62 /// If block has only one predecessor, and that predecessor has only one successor,
63 /// these two blocks can be fused as one.
64 fn fuse_block(&mut self, mut bb: BBPtr, func: FunPtr) -> Result<bool> {
65 let Some(mut pred) = bb.get_pred_bb().first().cloned() else {
66 return Ok(false);
67 };
68 if pred.get_succ_bb().len() == 1 && bb.get_pred_bb().len() == 1 {
69 // Last instruction is "br", move the rest to successor block
70 for inst in pred.iter_rev().skip(1) {
71 bb.push_front(inst);
72 }
73
74 // Replace `pred -> bb` with `bb`
75 // TODO remove requirement of func in replace_entry
76 pred.replace_entry(bb, func);
77
78 // Remove `pred`
79 pred.remove_self();
80 return Ok(true);
81 }
82 Ok(false)
83 }
84}