duskphantom_middle/transform/
store_elim.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::Result;
18
19use crate::{
20    analysis::memory_ssa::{MemorySSA, NodePtr},
21    ir::{instruction::InstType, InstPtr, Operand},
22    Program,
23};
24
25use super::Transform;
26
27pub fn optimize_program<'a>(
28    program: &'a mut Program,
29    memory_ssa: &'a mut MemorySSA<'a>,
30) -> Result<bool> {
31    StoreElim::new(program, memory_ssa).run_and_log()
32}
33
34pub struct StoreElim<'a> {
35    program: &'a mut Program,
36    memory_ssa: &'a mut MemorySSA<'a>,
37}
38
39impl<'a> Transform for StoreElim<'a> {
40    fn get_program_mut(&mut self) -> &mut Program {
41        self.program
42    }
43
44    fn name() -> String {
45        "store_elim".to_string()
46    }
47
48    fn run(&mut self) -> Result<bool> {
49        let mut changed = false;
50        for func in self.program.module.functions.clone() {
51            if func.is_main() {
52                for bb in func.po_iter() {
53                    for inst in bb.iter() {
54                        changed |= self.process_inst(inst)?;
55                    }
56                }
57            }
58        }
59        Ok(changed)
60    }
61}
62
63impl<'a> StoreElim<'a> {
64    pub fn new(program: &'a mut Program, memory_ssa: &'a mut MemorySSA<'a>) -> Self {
65        Self {
66            program,
67            memory_ssa,
68        }
69    }
70
71    /// Delete instruction and its corresponding MemorySSA node if it's not used.
72    /// This recurses into operands of the instruction.
73    fn process_inst(&mut self, inst: InstPtr) -> Result<bool> {
74        if !self.can_delete_inst(inst) {
75            return Ok(false);
76        }
77        if let Some(node) = self.memory_ssa.get_inst_node(inst) {
78            if !self.can_delete_node(node) {
79                return self.try_fuse_store(node);
80            }
81            self.remove_node(node)?;
82        };
83        self.remove_inst(inst)?;
84        Ok(true)
85    }
86
87    /// Delete MemorySSA node if unused.
88    /// This recurses into used nodes of the node.
89    fn process_node(&mut self, node: NodePtr) -> Result<bool> {
90        if !self.can_delete_node(node) {
91            return self.try_fuse_store(node);
92        }
93        if let Some(inst) = node.get_inst() {
94            if !self.can_delete_inst(inst) {
95                return Ok(false);
96            }
97            self.remove_inst(inst)?;
98        }
99        self.remove_node(node)?;
100        Ok(true)
101    }
102
103    /// Check if instruction can be deleted.
104    fn can_delete_inst(&self, inst: InstPtr) -> bool {
105        let no_io = !self.memory_ssa.effect_analysis.has_io(inst);
106        let no_user = inst.get_user().is_empty();
107        let no_control = !matches!(inst.get_type(), InstType::Br | InstType::Ret);
108        no_io && no_user && no_control
109    }
110
111    /// Check if node can be deleted.
112    fn can_delete_node(&self, node: NodePtr) -> bool {
113        self.memory_ssa.get_user(node).is_empty()
114    }
115
116    /// Remove instruction and recurse into operands.
117    fn remove_inst(&mut self, mut inst: InstPtr) -> Result<()> {
118        let operands: Vec<_> = inst.get_operand().into();
119        inst.remove_self();
120        for op in operands {
121            if let Operand::Instruction(inst) = op {
122                self.process_inst(inst)?;
123            }
124        }
125        Ok(())
126    }
127
128    /// Remove node and recurse into used nodes.
129    fn remove_node(&mut self, node: NodePtr) -> Result<()> {
130        let used_node = node.get_used_node();
131        self.memory_ssa.remove_node(node);
132        for node in used_node {
133            self.process_node(node)?;
134        }
135        Ok(())
136    }
137
138    /// Attempt to remove overridden store instructions, and fuse them into a single store.
139    fn try_fuse_store(&mut self, node: NodePtr) -> Result<bool> {
140        let Some(inst) = node.get_inst() else {
141            return Ok(false);
142        };
143
144        // Only explicit store can be stored
145        if inst.get_type() != InstType::Store {
146            return Ok(false);
147        }
148
149        // Get store position
150        let store_ptr = inst.get_operand()[1].clone();
151
152        // Traverse all stores upwards MemoryDef chain, get if it's used or overridden
153        let mut used = false;
154        let mut cursor = Some(node);
155        while let Some(curr) = cursor {
156            cursor = None;
157            let mut overridden = false;
158            for user in self.memory_ssa.get_user(curr).clone() {
159                if let Some(user_inst) = user.get_inst() {
160                    if user_inst.get_type() == InstType::Store {
161                        cursor = Some(user);
162                        let override_ptr = user_inst.get_operand()[1].clone();
163
164                        // For store, check if override
165                        if store_ptr == override_ptr {
166                            overridden = true;
167                        }
168                        continue;
169                    }
170                }
171
172                // For load / call / phi, mark the store as used
173                used = true;
174            }
175
176            // Break if used or overridden
177            if used || overridden {
178                break;
179            }
180        }
181
182        // Eliminate store if unused
183        if !used {
184            self.remove_inst(inst)?;
185            self.memory_ssa.remove_node(node);
186            return Ok(true);
187        }
188        Ok(false)
189    }
190}