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