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