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