1 //! Optimization driver using ISLE rewrite rules on an egraph.
2 
3 use crate::egraph::{NewOrExistingInst, OptimizeCtx};
4 pub use crate::ir::condcodes::{FloatCC, IntCC};
5 use crate::ir::dfg::ValueDef;
6 pub use crate::ir::immediates::{Ieee16, Ieee32, Ieee64, Ieee128, Imm64, Offset32, Uimm8, V128Imm};
7 use crate::ir::instructions::InstructionFormat;
8 pub use crate::ir::types::*;
9 pub use crate::ir::{
10     AtomicRmwOp, BlockCall, Constant, DynamicStackSlot, FuncRef, GlobalValue, Immediate,
11     InstructionData, MemFlags, Opcode, StackSlot, TrapCode, Type, Value,
12 };
13 use crate::isle_common_prelude_methods;
14 use crate::machinst::isle::*;
15 use crate::trace;
16 use cranelift_entity::packed_option::ReservedValue;
17 use smallvec::{SmallVec, smallvec};
18 use std::marker::PhantomData;
19 
20 #[allow(dead_code)]
21 pub type Unit = ();
22 pub type ValueArray2 = [Value; 2];
23 pub type ValueArray3 = [Value; 3];
24 
25 const MAX_ISLE_RETURNS: usize = 8;
26 
27 pub type ConstructorVec<T> = SmallVec<[T; MAX_ISLE_RETURNS]>;
28 
29 type TypeAndInstructionData = (Type, InstructionData);
30 
31 impl<T: smallvec::Array> generated_code::Length for SmallVec<T> {
32     #[inline]
33     fn len(&self) -> usize {
34         SmallVec::len(self)
35     }
36 }
37 
38 pub(crate) mod generated_code;
39 use generated_code::{ContextIter, IntoContextIter};
40 
41 pub(crate) struct IsleContext<'a, 'b, 'c> {
42     pub(crate) ctx: &'a mut OptimizeCtx<'b, 'c>,
43 }
44 
45 impl IsleContext<'_, '_, '_> {
46     pub(crate) fn dfg(&self) -> &crate::ir::DataFlowGraph {
47         &self.ctx.func.dfg
48     }
49 }
50 
51 pub(crate) struct InstDataEtorIter<'a, 'b, 'c> {
52     stack: SmallVec<[Value; 8]>,
53     _phantom1: PhantomData<&'a ()>,
54     _phantom2: PhantomData<&'b ()>,
55     _phantom3: PhantomData<&'c ()>,
56 }
57 
58 impl Default for InstDataEtorIter<'_, '_, '_> {
59     fn default() -> Self {
60         InstDataEtorIter {
61             stack: SmallVec::default(),
62             _phantom1: PhantomData,
63             _phantom2: PhantomData,
64             _phantom3: PhantomData,
65         }
66     }
67 }
68 
69 impl<'a, 'b, 'c> InstDataEtorIter<'a, 'b, 'c> {
70     fn new(root: Value) -> Self {
71         debug_assert_ne!(root, Value::reserved_value());
72         trace!("new iter from root {root}");
73         Self {
74             stack: smallvec![root],
75             _phantom1: PhantomData,
76             _phantom2: PhantomData,
77             _phantom3: PhantomData,
78         }
79     }
80 }
81 
82 impl<'a, 'b, 'c> ContextIter for InstDataEtorIter<'a, 'b, 'c>
83 where
84     'b: 'a,
85     'c: 'b,
86 {
87     type Context = IsleContext<'a, 'b, 'c>;
88     type Output = (Type, InstructionData);
89 
90     fn next(&mut self, ctx: &mut IsleContext<'a, 'b, 'c>) -> Option<Self::Output> {
91         while let Some(value) = self.stack.pop() {
92             debug_assert!(ctx.ctx.func.dfg.value_is_real(value));
93             trace!("iter: value {:?}", value);
94             match ctx.ctx.func.dfg.value_def(value) {
95                 ValueDef::Union(x, y) => {
96                     debug_assert_ne!(x, Value::reserved_value());
97                     debug_assert_ne!(y, Value::reserved_value());
98                     trace!(" -> {}, {}", x, y);
99                     self.stack.push(x);
100                     self.stack.push(y);
101                     continue;
102                 }
103                 ValueDef::Result(inst, _) if ctx.ctx.func.dfg.inst_results(inst).len() == 1 => {
104                     let ty = ctx.ctx.func.dfg.value_type(value);
105                     trace!(" -> value of type {}", ty);
106                     return Some((ty, ctx.ctx.func.dfg.insts[inst]));
107                 }
108                 _ => {}
109             }
110         }
111         None
112     }
113 }
114 
115 impl<'a, 'b, 'c> IntoContextIter for InstDataEtorIter<'a, 'b, 'c>
116 where
117     'b: 'a,
118     'c: 'b,
119 {
120     type Context = IsleContext<'a, 'b, 'c>;
121     type Output = (Type, InstructionData);
122     type IntoIter = Self;
123 
124     fn into_context_iter(self) -> Self {
125         self
126     }
127 }
128 
129 #[derive(Default)]
130 pub(crate) struct MaybeUnaryEtorIter<'a, 'b, 'c> {
131     opcode: Option<Opcode>,
132     inner: InstDataEtorIter<'a, 'b, 'c>,
133     fallback: Option<Value>,
134 }
135 
136 impl MaybeUnaryEtorIter<'_, '_, '_> {
137     fn new(opcode: Opcode, value: Value) -> Self {
138         debug_assert_eq!(opcode.format(), InstructionFormat::Unary);
139         Self {
140             opcode: Some(opcode),
141             inner: InstDataEtorIter::new(value),
142             fallback: Some(value),
143         }
144     }
145 }
146 
147 impl<'a, 'b, 'c> ContextIter for MaybeUnaryEtorIter<'a, 'b, 'c>
148 where
149     'b: 'a,
150     'c: 'b,
151 {
152     type Context = IsleContext<'a, 'b, 'c>;
153     type Output = (Type, Value);
154 
155     fn next(&mut self, ctx: &mut IsleContext<'a, 'b, 'c>) -> Option<Self::Output> {
156         debug_assert_ne!(self.opcode, None);
157         while let Some((ty, inst_def)) = self.inner.next(ctx) {
158             let InstructionData::Unary { opcode, arg } = inst_def else {
159                 continue;
160             };
161             if Some(opcode) == self.opcode {
162                 self.fallback = None;
163                 return Some((ty, arg));
164             }
165         }
166 
167         self.fallback.take().map(|value| {
168             let ty = generated_code::Context::value_type(ctx, value);
169             (ty, value)
170         })
171     }
172 }
173 
174 impl<'a, 'b, 'c> IntoContextIter for MaybeUnaryEtorIter<'a, 'b, 'c>
175 where
176     'b: 'a,
177     'c: 'b,
178 {
179     type Context = IsleContext<'a, 'b, 'c>;
180     type Output = (Type, Value);
181     type IntoIter = Self;
182 
183     fn into_context_iter(self) -> Self {
184         self
185     }
186 }
187 
188 impl<'a, 'b, 'c> generated_code::Context for IsleContext<'a, 'b, 'c> {
189     isle_common_prelude_methods!();
190 
191     type inst_data_value_etor_returns = InstDataEtorIter<'a, 'b, 'c>;
192 
193     fn inst_data_value_etor(&mut self, eclass: Value, returns: &mut InstDataEtorIter<'a, 'b, 'c>) {
194         *returns = InstDataEtorIter::new(eclass);
195     }
196 
197     type inst_data_value_tupled_etor_returns = InstDataEtorIter<'a, 'b, 'c>;
198 
199     fn inst_data_value_tupled_etor(
200         &mut self,
201         eclass: Value,
202         returns: &mut InstDataEtorIter<'a, 'b, 'c>,
203     ) {
204         // Literally identical to `inst_data_value_etor`, just a different nominal type in ISLE
205         self.inst_data_value_etor(eclass, returns);
206     }
207 
208     fn make_inst_ctor(&mut self, ty: Type, op: &InstructionData) -> Value {
209         trace!("make_inst_ctor: creating {:?}", op);
210         let value = self.ctx.insert_pure_enode(NewOrExistingInst::New(*op, ty));
211         trace!("make_inst_ctor: {:?} -> {}", op, value);
212         value
213     }
214 
215     fn make_skeleton_inst_ctor(&mut self, data: &InstructionData) -> Inst {
216         let inst = self.ctx.func.dfg.make_inst(*data);
217         self.ctx
218             .func
219             .dfg
220             .make_inst_results(inst, Default::default());
221         inst
222     }
223 
224     fn inst_data_etor(&mut self, inst: Inst) -> Option<InstructionData> {
225         Some(self.ctx.func.dfg.insts[inst])
226     }
227 
228     fn value_array_2_ctor(&mut self, arg0: Value, arg1: Value) -> ValueArray2 {
229         [arg0, arg1]
230     }
231 
232     fn value_array_3_ctor(&mut self, arg0: Value, arg1: Value, arg2: Value) -> ValueArray3 {
233         [arg0, arg1, arg2]
234     }
235 
236     #[inline]
237     fn value_type(&mut self, val: Value) -> Type {
238         self.ctx.func.dfg.value_type(val)
239     }
240 
241     fn iconst_sextend_etor(
242         &mut self,
243         (ty, inst_data): (Type, InstructionData),
244     ) -> Option<(Type, i64)> {
245         if let InstructionData::UnaryImm {
246             opcode: Opcode::Iconst,
247             imm,
248         } = inst_data
249         {
250             Some((ty, self.i64_sextend_imm64(ty, imm)))
251         } else {
252             None
253         }
254     }
255 
256     fn remat(&mut self, value: Value) -> Value {
257         trace!("remat: {}", value);
258         self.ctx.remat_values.insert(value);
259         self.ctx.stats.remat += 1;
260         value
261     }
262 
263     fn subsume(&mut self, value: Value) -> Value {
264         trace!("subsume: {}", value);
265         self.ctx.subsume_values.insert(value);
266         self.ctx.stats.subsume += 1;
267         value
268     }
269 
270     fn splat64(&mut self, val: u64) -> Constant {
271         let val = u128::from(val);
272         let val = val | (val << 64);
273         let imm = V128Imm(val.to_le_bytes());
274         self.ctx.func.dfg.constants.insert(imm.into())
275     }
276 
277     type sextend_maybe_etor_returns = MaybeUnaryEtorIter<'a, 'b, 'c>;
278     fn sextend_maybe_etor(&mut self, value: Value, returns: &mut Self::sextend_maybe_etor_returns) {
279         *returns = MaybeUnaryEtorIter::new(Opcode::Sextend, value);
280     }
281 
282     type uextend_maybe_etor_returns = MaybeUnaryEtorIter<'a, 'b, 'c>;
283     fn uextend_maybe_etor(&mut self, value: Value, returns: &mut Self::uextend_maybe_etor_returns) {
284         *returns = MaybeUnaryEtorIter::new(Opcode::Uextend, value);
285     }
286 
287     // NB: Cranelift's defined semantics for `fcvt_from_{s,u}int` match Rust's
288     // own semantics for converting an integer to a float, so these are all
289     // implemented with `as` conversions in Rust.
290     fn f32_from_uint(&mut self, n: u64) -> Ieee32 {
291         Ieee32::with_float(n as f32)
292     }
293 
294     fn f64_from_uint(&mut self, n: u64) -> Ieee64 {
295         Ieee64::with_float(n as f64)
296     }
297 
298     fn f32_from_sint(&mut self, n: i64) -> Ieee32 {
299         Ieee32::with_float(n as f32)
300     }
301 
302     fn f64_from_sint(&mut self, n: i64) -> Ieee64 {
303         Ieee64::with_float(n as f64)
304     }
305 
306     fn u64_bswap16(&mut self, n: u64) -> u64 {
307         (n as u16).swap_bytes() as u64
308     }
309 
310     fn u64_bswap32(&mut self, n: u64) -> u64 {
311         (n as u32).swap_bytes() as u64
312     }
313 
314     fn u64_bswap64(&mut self, n: u64) -> u64 {
315         n.swap_bytes()
316     }
317 
318     fn ieee128_constant_extractor(&mut self, n: Constant) -> Option<Ieee128> {
319         self.ctx.func.dfg.constants.get(n).try_into().ok()
320     }
321 
322     fn ieee128_constant(&mut self, n: Ieee128) -> Constant {
323         self.ctx.func.dfg.constants.insert(n.into())
324     }
325 }
326