Skip to main content

cranelift_codegen/
opts.rs

1//! Optimization driver using ISLE rewrite rules on an egraph.
2
3mod div_const;
4
5use crate::egraph::{NewOrExistingInst, OptimizeCtx};
6pub use crate::ir::condcodes::{FloatCC, IntCC};
7use crate::ir::dfg::ValueDef;
8pub use crate::ir::immediates::{Ieee16, Ieee32, Ieee64, Ieee128, Imm64, Offset32, Uimm8, V128Imm};
9use crate::ir::instructions::InstructionFormat;
10pub use crate::ir::types::*;
11pub use crate::ir::{
12    AtomicRmwOp, Block, BlockCall, Constant, DynamicStackSlot, FuncRef, GlobalValue, Immediate,
13    InstructionData, JumpTable, MemFlagsData, Opcode, StackSlot, TrapCode, Type, Value,
14};
15use crate::isle_common_prelude_methods;
16use crate::machinst::isle::*;
17use crate::trace;
18use core::marker::PhantomData;
19use cranelift_entity::packed_option::ReservedValue;
20use smallvec::{SmallVec, smallvec};
21
22pub type Unit = ();
23pub type ValueArray2 = [Value; 2];
24pub type ValueArray3 = [Value; 3];
25
26const MAX_ISLE_RETURNS: usize = 8;
27
28pub type ConstructorVec<T> = SmallVec<[T; MAX_ISLE_RETURNS]>;
29
30type TypeAndInstructionData = (Type, InstructionData);
31
32impl<T: smallvec::Array> generated_code::Length for SmallVec<T> {
33    #[inline]
34    fn len(&self) -> usize {
35        SmallVec::len(self)
36    }
37}
38
39pub(crate) mod generated_code;
40use generated_code::{ContextIter, IntoContextIter};
41
42pub(crate) struct IsleContext<'a, 'b, 'c> {
43    pub(crate) ctx: &'a mut OptimizeCtx<'b, 'c>,
44}
45
46impl IsleContext<'_, '_, '_> {
47    #[allow(dead_code, reason = "dead code, only on nightly rust at this time")]
48    pub(crate) fn dfg(&self) -> &crate::ir::DataFlowGraph {
49        &self.ctx.func.dfg
50    }
51}
52
53pub(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
60impl Default for InstDataEtorIter<'_, '_, '_> {
61    fn default() -> Self {
62        InstDataEtorIter {
63            stack: SmallVec::default(),
64            _phantom1: PhantomData,
65            _phantom2: PhantomData,
66            _phantom3: PhantomData,
67        }
68    }
69}
70
71impl<'a, 'b, 'c> InstDataEtorIter<'a, 'b, 'c> {
72    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
84impl<'a, 'b, 'c> ContextIter for InstDataEtorIter<'a, 'b, 'c>
85where
86    'b: 'a,
87    'c: 'b,
88{
89    type Context = IsleContext<'a, 'b, 'c>;
90    type Output = (Type, InstructionData);
91
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                    // Charge one unit of fuel per yielded match. When
107                    // fuel is exhausted, terminate iteration early:
108                    // returning no matches is always semantically valid
109                    // (we just skip would-be rewrites) and bounds work
110                    // per top-level ISLE invocation.
111                    if ctx.ctx.extractor_fuel == 0 {
112                        ctx.ctx.stats.rewrite_fuel_exhausted += 1;
113                        trace!(" -> rewrite fuel exhausted");
114                        return None;
115                    }
116                    ctx.ctx.extractor_fuel -= 1;
117                    let ty = ctx.ctx.func.dfg.value_type(value);
118                    trace!(" -> value of type {}", ty);
119                    return Some((ty, ctx.ctx.func.dfg.insts[inst]));
120                }
121                _ => {}
122            }
123        }
124        None
125    }
126}
127
128impl<'a, 'b, 'c> IntoContextIter for InstDataEtorIter<'a, 'b, 'c>
129where
130    'b: 'a,
131    'c: 'b,
132{
133    type Context = IsleContext<'a, 'b, 'c>;
134    type Output = (Type, InstructionData);
135    type IntoIter = Self;
136
137    fn into_context_iter(self) -> Self {
138        self
139    }
140}
141
142#[derive(Default)]
143pub(crate) struct MaybeUnaryEtorIter<'a, 'b, 'c> {
144    opcode: Option<Opcode>,
145    inner: InstDataEtorIter<'a, 'b, 'c>,
146    fallback: Option<Value>,
147}
148
149impl MaybeUnaryEtorIter<'_, '_, '_> {
150    fn new(opcode: Opcode, value: Value) -> Self {
151        debug_assert_eq!(opcode.format(), InstructionFormat::Unary);
152        Self {
153            opcode: Some(opcode),
154            inner: InstDataEtorIter::new(value),
155            fallback: Some(value),
156        }
157    }
158}
159
160impl<'a, 'b, 'c> ContextIter for MaybeUnaryEtorIter<'a, 'b, 'c>
161where
162    'b: 'a,
163    'c: 'b,
164{
165    type Context = IsleContext<'a, 'b, 'c>;
166    type Output = (Type, Value);
167
168    fn next(&mut self, ctx: &mut IsleContext<'a, 'b, 'c>) -> Option<Self::Output> {
169        debug_assert_ne!(self.opcode, None);
170        while let Some((ty, inst_def)) = self.inner.next(ctx) {
171            let InstructionData::Unary { opcode, arg } = inst_def else {
172                continue;
173            };
174            if Some(opcode) == self.opcode {
175                self.fallback = None;
176                return Some((ty, arg));
177            }
178        }
179
180        self.fallback.take().map(|value| {
181            let ty = generated_code::Context::value_type(ctx, value);
182            (ty, value)
183        })
184    }
185}
186
187impl<'a, 'b, 'c> IntoContextIter for MaybeUnaryEtorIter<'a, 'b, 'c>
188where
189    'b: 'a,
190    'c: 'b,
191{
192    type Context = IsleContext<'a, 'b, 'c>;
193    type Output = (Type, Value);
194    type IntoIter = Self;
195
196    fn into_context_iter(self) -> Self {
197        self
198    }
199}
200
201impl<'a, 'b, 'c> generated_code::Context for IsleContext<'a, 'b, 'c> {
202    isle_common_prelude_methods!();
203
204    fn zero_constant(&mut self, ty: Type) -> Constant {
205        let data = vec![0; ty.bytes() as usize];
206        self.ctx.func.dfg.constants.insert(data.into())
207    }
208
209    fn f16_zero(&mut self) -> Ieee16 {
210        Ieee16::with_bits(0)
211    }
212
213    fn ty_vector(&mut self, ty: Type) -> Option<Type> {
214        ty.is_vector().then_some(ty)
215    }
216
217    type inst_data_value_etor_returns = InstDataEtorIter<'a, 'b, 'c>;
218
219    fn inst_data_value_etor(&mut self, eclass: Value, returns: &mut InstDataEtorIter<'a, 'b, 'c>) {
220        *returns = InstDataEtorIter::new(eclass);
221    }
222
223    type inst_data_value_tupled_etor_returns = InstDataEtorIter<'a, 'b, 'c>;
224
225    fn inst_data_value_tupled_etor(
226        &mut self,
227        eclass: Value,
228        returns: &mut InstDataEtorIter<'a, 'b, 'c>,
229    ) {
230        // Literally identical to `inst_data_value_etor`, just a different nominal type in ISLE
231        self.inst_data_value_etor(eclass, returns);
232    }
233
234    fn make_inst_ctor(&mut self, ty: Type, op: &InstructionData) -> Value {
235        trace!("make_inst_ctor: creating {:?}", op);
236        let value = self.ctx.insert_pure_enode(NewOrExistingInst::New(*op, ty));
237        trace!("make_inst_ctor: {:?} -> {}", op, value);
238        value
239    }
240
241    fn make_skeleton_inst_ctor(&mut self, data: &InstructionData) -> Inst {
242        let inst = self.ctx.func.dfg.make_inst(*data);
243        self.ctx
244            .func
245            .dfg
246            .make_inst_results(inst, Default::default());
247        inst
248    }
249
250    fn inst_data_etor(&mut self, inst: Inst) -> Option<InstructionData> {
251        Some(self.ctx.func.dfg.insts[inst])
252    }
253
254    fn value_array_2_ctor(&mut self, arg0: Value, arg1: Value) -> ValueArray2 {
255        [arg0, arg1]
256    }
257
258    fn value_array_3_ctor(&mut self, arg0: Value, arg1: Value, arg2: Value) -> ValueArray3 {
259        [arg0, arg1, arg2]
260    }
261
262    #[inline]
263    fn value_type(&mut self, val: Value) -> Type {
264        self.ctx.func.dfg.value_type(val)
265    }
266
267    fn resolve_jump_table_entry(&mut self, table: JumpTable, index: u64) -> BlockCall {
268        let jt_data = &self.ctx.func.dfg.jump_tables[table];
269        let entries = jt_data.as_slice();
270        if let Ok(index) = usize::try_from(index)
271            && index < entries.len()
272        {
273            entries[index]
274        } else {
275            jt_data.default_block()
276        }
277    }
278
279    fn block_call_block(&mut self, block_call: BlockCall) -> Block {
280        block_call.block(&self.ctx.func.dfg.value_lists)
281    }
282
283    fn just_trap_block(&mut self, block: &Block) -> Option<TrapCode> {
284        self.ctx
285            .branch_to_trap_analysis
286            .analyze_block(self.ctx.func, *block)
287    }
288
289    fn iconst_sextend_etor(
290        &mut self,
291        (ty, inst_data): (Type, InstructionData),
292    ) -> Option<(Type, i64)> {
293        if let InstructionData::UnaryImm {
294            opcode: Opcode::Iconst,
295            imm,
296        } = inst_data
297        {
298            Some((ty, self.i64_sextend_imm64(ty, imm)))
299        } else {
300            None
301        }
302    }
303
304    fn all_zero_etor(&mut self, (ty, inst_data): (Type, InstructionData)) -> Option<Type> {
305        let is_all_zero = match inst_data {
306            InstructionData::UnaryImm {
307                opcode: Opcode::Iconst,
308                imm,
309            } => imm.bits() == 0,
310            InstructionData::UnaryIeee16 {
311                opcode: Opcode::F16const,
312                imm,
313            } => imm.bits() == 0,
314            InstructionData::UnaryIeee32 {
315                opcode: Opcode::F32const,
316                imm,
317            } => imm.bits() == 0,
318            InstructionData::UnaryIeee64 {
319                opcode: Opcode::F64const,
320                imm,
321            } => imm.bits() == 0,
322            InstructionData::UnaryConst {
323                opcode: Opcode::F128const | Opcode::Vconst,
324                constant_handle,
325            } => {
326                let constant = self.ctx.func.dfg.constants.get(constant_handle);
327                constant.len() == ty.bytes() as usize && constant.iter().all(|&byte| byte == 0)
328            }
329            _ => false,
330        };
331
332        is_all_zero.then_some(ty)
333    }
334
335    fn remat(&mut self, value: Value) -> Value {
336        trace!("remat: {}", value);
337        self.ctx.remat_values.insert(value);
338        self.ctx.stats.remat += 1;
339        value
340    }
341
342    fn subsume(&mut self, value: Value) -> Value {
343        trace!("subsume: {}", value);
344        self.ctx.subsume_values.insert(value);
345        self.ctx.stats.subsume += 1;
346        value
347    }
348
349    fn splat64(&mut self, val: u64) -> Constant {
350        let val = u128::from(val);
351        let val = val | (val << 64);
352        let imm = V128Imm(val.to_le_bytes());
353        self.ctx.func.dfg.constants.insert(imm.into())
354    }
355
356    fn scalar_to_vector_const64(&mut self, val: u64) -> Constant {
357        let imm = V128Imm(u128::from(val).to_le_bytes());
358        self.ctx.func.dfg.constants.insert(imm.into())
359    }
360
361    type sextend_maybe_etor_returns = MaybeUnaryEtorIter<'a, 'b, 'c>;
362    fn sextend_maybe_etor(&mut self, value: Value, returns: &mut Self::sextend_maybe_etor_returns) {
363        *returns = MaybeUnaryEtorIter::new(Opcode::Sextend, value);
364    }
365
366    type uextend_maybe_etor_returns = MaybeUnaryEtorIter<'a, 'b, 'c>;
367    fn uextend_maybe_etor(&mut self, value: Value, returns: &mut Self::uextend_maybe_etor_returns) {
368        *returns = MaybeUnaryEtorIter::new(Opcode::Uextend, value);
369    }
370
371    // NB: Cranelift's defined semantics for `fcvt_from_{s,u}int` match Rust's
372    // own semantics for converting an integer to a float, so these are all
373    // implemented with `as` conversions in Rust.
374    fn f32_from_uint(&mut self, n: u64) -> Ieee32 {
375        Ieee32::with_float(n as f32)
376    }
377
378    fn f64_from_uint(&mut self, n: u64) -> Ieee64 {
379        Ieee64::with_float(n as f64)
380    }
381
382    fn f32_from_sint(&mut self, n: i64) -> Ieee32 {
383        Ieee32::with_float(n as f32)
384    }
385
386    fn f64_from_sint(&mut self, n: i64) -> Ieee64 {
387        Ieee64::with_float(n as f64)
388    }
389
390    fn u64_bswap16(&mut self, n: u64) -> u64 {
391        (n as u16).swap_bytes() as u64
392    }
393
394    fn u64_bswap32(&mut self, n: u64) -> u64 {
395        (n as u32).swap_bytes() as u64
396    }
397
398    fn u64_bswap64(&mut self, n: u64) -> u64 {
399        n.swap_bytes()
400    }
401
402    fn ieee128_constant_extractor(&mut self, n: Constant) -> Option<Ieee128> {
403        self.ctx.func.dfg.constants.get(n).try_into().ok()
404    }
405
406    fn ieee128_constant(&mut self, n: Ieee128) -> Constant {
407        self.ctx.func.dfg.constants.insert(n.into())
408    }
409
410    fn div_const_magic_u32(&mut self, d: u32) -> generated_code::DivConstMagicU32 {
411        let div_const::MU32 {
412            mul_by,
413            do_add,
414            shift_by,
415        } = div_const::magic_u32(d);
416        generated_code::DivConstMagicU32::U32 {
417            mul_by,
418            do_add,
419            shift_by: shift_by.try_into().unwrap(),
420        }
421    }
422
423    fn div_const_magic_u64(&mut self, d: u64) -> generated_code::DivConstMagicU64 {
424        let div_const::MU64 {
425            mul_by,
426            do_add,
427            shift_by,
428        } = div_const::magic_u64(d);
429        generated_code::DivConstMagicU64::U64 {
430            mul_by,
431            do_add,
432            shift_by: shift_by.try_into().unwrap(),
433        }
434    }
435
436    fn div_const_magic_s32(&mut self, d: i32) -> generated_code::DivConstMagicS32 {
437        let div_const::MS32 { mul_by, shift_by } = div_const::magic_s32(d);
438        generated_code::DivConstMagicS32::S32 {
439            mul_by,
440            shift_by: shift_by.try_into().unwrap(),
441        }
442    }
443
444    fn div_const_magic_s64(&mut self, d: i64) -> generated_code::DivConstMagicS64 {
445        let div_const::MS64 { mul_by, shift_by } = div_const::magic_s64(d);
446        generated_code::DivConstMagicS64::S64 {
447            mul_by,
448            shift_by: shift_by.try_into().unwrap(),
449        }
450    }
451}