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::{Ieee128, Ieee16, Ieee32, Ieee64, 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 Range = (usize, usize);
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]
34     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 pub(crate) struct InstDataEtorIter<'a, 'b, 'c> {
47     stack: SmallVec<[Value; 8]>,
48     _phantom1: PhantomData<&'a ()>,
49     _phantom2: PhantomData<&'b ()>,
50     _phantom3: PhantomData<&'c ()>,
51 }
52 
53 impl Default for InstDataEtorIter<'_, '_, '_> {
54     fn default() -> Self {
55         InstDataEtorIter {
56             stack: SmallVec::default(),
57             _phantom1: PhantomData,
58             _phantom2: PhantomData,
59             _phantom3: PhantomData,
60         }
61     }
62 }
63 
64 impl<'a, 'b, 'c> InstDataEtorIter<'a, 'b, 'c> {
65     fn new(root: Value) -> Self {
66         debug_assert_ne!(root, Value::reserved_value());
67         Self {
68             stack: smallvec![root],
69             _phantom1: PhantomData,
70             _phantom2: PhantomData,
71             _phantom3: PhantomData,
72         }
73     }
74 }
75 
76 impl<'a, 'b, 'c> ContextIter for InstDataEtorIter<'a, 'b, 'c>
77 where
78     'b: 'a,
79     'c: 'b,
80 {
81     type Context = IsleContext<'a, 'b, 'c>;
82     type Output = (Type, InstructionData);
83 
84     fn next(&mut self, ctx: &mut IsleContext<'a, 'b, 'c>) -> Option<Self::Output> {
85         while let Some(value) = self.stack.pop() {
86             debug_assert!(ctx.ctx.func.dfg.value_is_real(value));
87             trace!("iter: value {:?}", value);
88             match ctx.ctx.func.dfg.value_def(value) {
89                 ValueDef::Union(x, y) => {
90                     debug_assert_ne!(x, Value::reserved_value());
91                     debug_assert_ne!(y, Value::reserved_value());
92                     trace!(" -> {}, {}", x, y);
93                     self.stack.push(x);
94                     self.stack.push(y);
95                     continue;
96                 }
97                 ValueDef::Result(inst, _) if ctx.ctx.func.dfg.inst_results(inst).len() == 1 => {
98                     let ty = ctx.ctx.func.dfg.value_type(value);
99                     trace!(" -> value of type {}", ty);
100                     return Some((ty, ctx.ctx.func.dfg.insts[inst]));
101                 }
102                 _ => {}
103             }
104         }
105         None
106     }
107 }
108 
109 impl<'a, 'b, 'c> IntoContextIter for InstDataEtorIter<'a, 'b, 'c>
110 where
111     'b: 'a,
112     'c: 'b,
113 {
114     type Context = IsleContext<'a, 'b, 'c>;
115     type Output = (Type, InstructionData);
116     type IntoIter = Self;
117 
118     fn into_context_iter(self) -> Self {
119         self
120     }
121 }
122 
123 #[derive(Default)]
124 pub(crate) struct MaybeUnaryEtorIter<'a, 'b, 'c> {
125     opcode: Option<Opcode>,
126     inner: InstDataEtorIter<'a, 'b, 'c>,
127     fallback: Option<Value>,
128 }
129 
130 impl MaybeUnaryEtorIter<'_, '_, '_> {
131     fn new(opcode: Opcode, value: Value) -> Self {
132         debug_assert_eq!(opcode.format(), InstructionFormat::Unary);
133         Self {
134             opcode: Some(opcode),
135             inner: InstDataEtorIter::new(value),
136             fallback: Some(value),
137         }
138     }
139 }
140 
141 impl<'a, 'b, 'c> ContextIter for MaybeUnaryEtorIter<'a, 'b, 'c>
142 where
143     'b: 'a,
144     'c: 'b,
145 {
146     type Context = IsleContext<'a, 'b, 'c>;
147     type Output = (Type, Value);
148 
149     fn next(&mut self, ctx: &mut IsleContext<'a, 'b, 'c>) -> Option<Self::Output> {
150         debug_assert_ne!(self.opcode, None);
151         while let Some((ty, inst_def)) = self.inner.next(ctx) {
152             let InstructionData::Unary { opcode, arg } = inst_def else {
153                 continue;
154             };
155             if Some(opcode) == self.opcode {
156                 self.fallback = None;
157                 return Some((ty, arg));
158             }
159         }
160 
161         self.fallback.take().map(|value| {
162             let ty = generated_code::Context::value_type(ctx, value);
163             (ty, value)
164         })
165     }
166 }
167 
168 impl<'a, 'b, 'c> IntoContextIter for MaybeUnaryEtorIter<'a, 'b, 'c>
169 where
170     'b: 'a,
171     'c: 'b,
172 {
173     type Context = IsleContext<'a, 'b, 'c>;
174     type Output = (Type, Value);
175     type IntoIter = Self;
176 
177     fn into_context_iter(self) -> Self {
178         self
179     }
180 }
181 
182 impl<'a, 'b, 'c> generated_code::Context for IsleContext<'a, 'b, 'c> {
183     isle_common_prelude_methods!();
184 
185     type inst_data_etor_returns = InstDataEtorIter<'a, 'b, 'c>;
186 
187     fn inst_data_etor(&mut self, eclass: Value, returns: &mut InstDataEtorIter<'a, 'b, 'c>) {
188         *returns = InstDataEtorIter::new(eclass);
189     }
190 
191     type inst_data_tupled_etor_returns = InstDataEtorIter<'a, 'b, 'c>;
192 
193     fn inst_data_tupled_etor(&mut self, eclass: Value, returns: &mut InstDataEtorIter<'a, 'b, 'c>) {
194         // Literally identical to `inst_data_etor`, just a different nominal type in ISLE
195         self.inst_data_etor(eclass, returns);
196     }
197 
198     fn make_inst_ctor(&mut self, ty: Type, op: &InstructionData) -> Value {
199         let value = self.ctx.insert_pure_enode(NewOrExistingInst::New(*op, ty));
200         trace!("make_inst_ctor: {:?} -> {}", op, value);
201         value
202     }
203 
204     fn value_array_2_ctor(&mut self, arg0: Value, arg1: Value) -> ValueArray2 {
205         [arg0, arg1]
206     }
207 
208     fn value_array_3_ctor(&mut self, arg0: Value, arg1: Value, arg2: Value) -> ValueArray3 {
209         [arg0, arg1, arg2]
210     }
211 
212     #[inline]
213     fn value_type(&mut self, val: Value) -> Type {
214         self.ctx.func.dfg.value_type(val)
215     }
216 
217     fn iconst_sextend_etor(
218         &mut self,
219         (ty, inst_data): (Type, InstructionData),
220     ) -> Option<(Type, i64)> {
221         if let InstructionData::UnaryImm {
222             opcode: Opcode::Iconst,
223             imm,
224         } = inst_data
225         {
226             Some((ty, self.i64_sextend_imm64(ty, imm)))
227         } else {
228             None
229         }
230     }
231 
232     fn remat(&mut self, value: Value) -> Value {
233         trace!("remat: {}", value);
234         self.ctx.remat_values.insert(value);
235         self.ctx.stats.remat += 1;
236         value
237     }
238 
239     fn subsume(&mut self, value: Value) -> Value {
240         trace!("subsume: {}", value);
241         self.ctx.subsume_values.insert(value);
242         self.ctx.stats.subsume += 1;
243         value
244     }
245 
246     fn splat64(&mut self, val: u64) -> Constant {
247         let val = u128::from(val);
248         let val = val | (val << 64);
249         let imm = V128Imm(val.to_le_bytes());
250         self.ctx.func.dfg.constants.insert(imm.into())
251     }
252 
253     type sextend_maybe_etor_returns = MaybeUnaryEtorIter<'a, 'b, 'c>;
254     fn sextend_maybe_etor(&mut self, value: Value, returns: &mut Self::sextend_maybe_etor_returns) {
255         *returns = MaybeUnaryEtorIter::new(Opcode::Sextend, value);
256     }
257 
258     type uextend_maybe_etor_returns = MaybeUnaryEtorIter<'a, 'b, 'c>;
259     fn uextend_maybe_etor(&mut self, value: Value, returns: &mut Self::uextend_maybe_etor_returns) {
260         *returns = MaybeUnaryEtorIter::new(Opcode::Uextend, value);
261     }
262 
263     // NB: Cranelift's defined semantics for `fcvt_from_{s,u}int` match Rust's
264     // own semantics for converting an integer to a float, so these are all
265     // implemented with `as` conversions in Rust.
266     fn f32_from_uint(&mut self, n: u64) -> Ieee32 {
267         Ieee32::with_float(n as f32)
268     }
269 
270     fn f64_from_uint(&mut self, n: u64) -> Ieee64 {
271         Ieee64::with_float(n as f64)
272     }
273 
274     fn f32_from_sint(&mut self, n: i64) -> Ieee32 {
275         Ieee32::with_float(n as f32)
276     }
277 
278     fn f64_from_sint(&mut self, n: i64) -> Ieee64 {
279         Ieee64::with_float(n as f64)
280     }
281 
282     fn u64_bswap16(&mut self, n: u64) -> u64 {
283         (n as u16).swap_bytes() as u64
284     }
285 
286     fn u64_bswap32(&mut self, n: u64) -> u64 {
287         (n as u32).swap_bytes() as u64
288     }
289 
290     fn u64_bswap64(&mut self, n: u64) -> u64 {
291         n.swap_bytes()
292     }
293 
294     fn ieee128_constant_extractor(&mut self, n: Constant) -> Option<Ieee128> {
295         self.ctx.func.dfg.constants.get(n).try_into().ok()
296     }
297 
298     fn ieee128_constant(&mut self, n: Ieee128) -> Constant {
299         self.ctx.func.dfg.constants.insert(n.into())
300     }
301 }
302