llvm_transforms/
ipcp.rs

1//! Inter-procedural constant propagation (IPCP).
2
3use crate::{const_prop::ConstProp, pass::ModulePass, value_rewrite::rewrite_values_in_kind};
4use llvm_ir::{
5    ArgId, Context, Function, FunctionId, GlobalId, InstrKind, Module, ValueRef,
6};
7
8/// Simple IPCP pass:
9/// - detect a direct callee argument that is constant across all direct callsites
10/// - clone callee and substitute that argument with the constant
11/// - redirect matching callsites to the specialized clone
12pub struct Ipcp;
13
14impl ModulePass for Ipcp {
15    fn name(&self) -> &'static str {
16        "ipcp"
17    }
18
19    fn run_on_module(&mut self, ctx: &mut Context, module: &mut Module) -> bool {
20        let mut changed = false;
21
22        for callee_idx in 0..module.functions.len() {
23            let callee_id = FunctionId(callee_idx as u32);
24            if module.functions[callee_idx].is_declaration || module.functions[callee_idx].args.is_empty() {
25                continue;
26            }
27
28            let callsites = collect_direct_calls(module, callee_id);
29            if callsites.is_empty() {
30                continue;
31            }
32            let Some((arg_idx, const_val)) = find_constant_arg(&callsites) else {
33                continue;
34            };
35
36            let spec_name = format!(
37                "{}.ipcp.a{}.c{}",
38                module.functions[callee_idx].name, arg_idx, const_val.0
39            );
40            let spec_id = if let Some(fid) = module.get_function_id(&spec_name) {
41                fid
42            } else {
43                let mut spec = clone_specialized_function(
44                    &module.functions[callee_idx],
45                    &spec_name,
46                    ArgId(arg_idx as u32),
47                    ValueRef::Constant(const_val),
48                );
49                let mut cp = ConstProp;
50                let _ = crate::pass::FunctionPass::run_on_function(&mut cp, ctx, &mut spec);
51                module.add_function(spec)
52            };
53
54            let spec_ty = module.functions[spec_id.0 as usize].ty;
55            for cs in callsites {
56                if cs.const_args.get(arg_idx).copied() == Some(Some(const_val)) {
57                    let instr = &mut module.functions[cs.caller.0 as usize].instructions[cs.iid.0 as usize];
58                    if let InstrKind::Call { callee, callee_ty, .. } = &mut instr.kind {
59                        *callee = ValueRef::Global(GlobalId(spec_id.0));
60                        *callee_ty = spec_ty;
61                        changed = true;
62                    }
63                }
64            }
65        }
66
67        changed
68    }
69}
70
71#[derive(Clone)]
72struct DirectCallSite {
73    caller: FunctionId,
74    iid: llvm_ir::InstrId,
75    const_args: Vec<Option<llvm_ir::ConstId>>,
76}
77
78fn collect_direct_calls(module: &Module, callee_id: FunctionId) -> Vec<DirectCallSite> {
79    let mut out = Vec::new();
80    for (caller_idx, f) in module.functions.iter().enumerate() {
81        if f.is_declaration {
82            continue;
83        }
84        for (iid_idx, instr) in f.instructions.iter().enumerate() {
85            let InstrKind::Call { callee, args, .. } = &instr.kind else {
86                continue;
87            };
88            if *callee != ValueRef::Global(GlobalId(callee_id.0)) {
89                continue;
90            }
91            let const_args = args
92                .iter()
93                .map(|a| match a {
94                    ValueRef::Constant(c) => Some(*c),
95                    _ => None,
96                })
97                .collect();
98            out.push(DirectCallSite {
99                caller: FunctionId(caller_idx as u32),
100                iid: llvm_ir::InstrId(iid_idx as u32),
101                const_args,
102            });
103        }
104    }
105    out
106}
107
108fn find_constant_arg(callsites: &[DirectCallSite]) -> Option<(usize, llvm_ir::ConstId)> {
109    let argc = callsites.first()?.const_args.len();
110    for ai in 0..argc {
111        let mut cst: Option<llvm_ir::ConstId> = None;
112        let mut ok = true;
113        for cs in callsites {
114            match cs.const_args.get(ai).copied().flatten() {
115                Some(c) => {
116                    if let Some(prev) = cst {
117                        if prev != c {
118                            ok = false;
119                            break;
120                        }
121                    } else {
122                        cst = Some(c);
123                    }
124                }
125                None => {
126                    ok = false;
127                    break;
128                }
129            }
130        }
131        if ok {
132            return cst.map(|c| (ai, c));
133        }
134    }
135    None
136}
137
138fn clone_specialized_function(
139    src: &Function,
140    new_name: &str,
141    arg_id: ArgId,
142    const_val: ValueRef,
143) -> Function {
144    let mut dst = Function::new(new_name.to_string(), src.ty, src.args.clone(), src.linkage);
145    dst.blocks = src.blocks.clone();
146    dst.instructions = src.instructions.clone();
147    dst.value_names = src.value_names.clone();
148    dst.arg_names = src.arg_names.clone();
149    dst.is_declaration = false;
150
151    for instr in &mut dst.instructions {
152        let old = instr.kind.clone();
153        instr.kind = rewrite_values_in_kind(old, |v| {
154            if v == ValueRef::Argument(arg_id) {
155                const_val
156            } else {
157                v
158            }
159        });
160    }
161    dst
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167    use llvm_ir::{Builder, Linkage};
168
169    #[test]
170    fn ipcp_specializes_constant_argument_and_rewrites_calls() {
171        let mut ctx = Context::new();
172        let mut module = Module::new("m");
173        let mut b = Builder::new(&mut ctx, &mut module);
174
175        b.add_function(
176            "addk",
177            b.ctx.i64_ty,
178            vec![b.ctx.i64_ty, b.ctx.i64_ty],
179            vec!["x".into(), "k".into()],
180            false,
181            Linkage::External,
182        );
183        let addk_entry = b.add_block("addk.entry");
184        b.position_at_end(addk_entry);
185        let x = b.get_arg(0);
186        let k = b.get_arg(1);
187        let s = b.build_add("s", x, k);
188        b.build_ret(s);
189
190        b.add_function(
191            "caller",
192            b.ctx.i64_ty,
193            vec![b.ctx.i64_ty],
194            vec!["x".into()],
195            false,
196            Linkage::External,
197        );
198        let caller_entry = b.add_block("caller.entry");
199        b.position_at_end(caller_entry);
200        let x0 = b.get_arg(0);
201        let c7 = b.const_int(b.ctx.i64_ty, 7);
202        let c7b = b.const_int(b.ctx.i64_ty, 7);
203        let call_ty = b.ctx.mk_fn_type(b.ctx.i64_ty, vec![b.ctx.i64_ty, b.ctx.i64_ty], false);
204        let t1 = b.build_call(
205            "t1",
206            b.ctx.i64_ty,
207            call_ty,
208            ValueRef::Global(GlobalId(0)),
209            vec![x0, c7],
210        );
211        let t2 = b.build_call(
212            "t2",
213            b.ctx.i64_ty,
214            call_ty,
215            ValueRef::Global(GlobalId(0)),
216            vec![t1, c7],
217        );
218        let _t3 = b.build_call(
219            "t3",
220            b.ctx.i64_ty,
221            call_ty,
222            ValueRef::Global(GlobalId(0)),
223            vec![t2, c7b],
224        );
225        b.build_ret(t2);
226
227        let mut pass = Ipcp;
228        let changed = pass.run_on_module(&mut ctx, &mut module);
229        assert!(changed);
230        assert!(
231            module.functions.iter().any(|f| f.name.starts_with("addk.ipcp")),
232            "expected specialized clone"
233        );
234    }
235}
llvm_transforms/ipcp.rs

llvm_transforms/
ipcp.rs
