pliron_llvm/
interface_impls.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) The pliron contributors
3
4//! Implementation of various op interfaces for LLVM IR instructions.
5
6use alloc::{boxed::Box, vec, vec::Vec};
7use core::num::NonZero;
8use thiserror::Error;
9
10use pliron::{
11    arg_err,
12    attribute::{AttrObj, Attribute, attr_cast},
13    basic_block::BasicBlock,
14    builtin::{
15        attr_interfaces::FloatAttr,
16        attributes::IntegerAttr,
17        op_interfaces::{BranchOpInterface, OneResultInterface},
18        types::{IntegerType, Signedness},
19    },
20    context::{Context, Ptr},
21    derive::op_interface_impl,
22    irbuild::{IRStatus, inserter::Inserter, rewriter::Rewriter},
23    op::Op,
24    opts::{
25        constants::{BranchOpFoldInterface, ConstFoldInterface},
26        dce::{BlockArgRemoval, SideEffects},
27        mem2reg::{
28            AllocInfo, PromotableAllocationInterface, PromotableOpInterface, PromotableOpKind,
29        },
30    },
31    result::Result,
32    utils::apint::{APInt, bw},
33    value::Value,
34};
35
36use crate::{
37    attributes::{FastmathFlags, FastmathFlagsAttr, ICmpPredicateAttr, IntegerOverflowFlagsAttr},
38    op_interfaces::{FastMathFlags, IntBinArithOpWithOverflowFlag, NNegFlag, PointerTypeResult},
39    ops::{
40        AShrOp, AddOp, AddressOfOp, AllocaOp, AndOp, BitcastOp, BrOp, CondBrOp, ConstantOp,
41        ExtractElementOp, ExtractValueOp, FAddOp, FCmpOp, FDivOp, FMulOp, FNegOp, FPExtOp,
42        FPToSIOp, FPToUIOp, FPTruncOp, FRemOp, FSubOp, FreezeOp, FuncOp, GetElementPtrOp, ICmpOp,
43        InsertElementOp, InsertValueOp, IntToPtrOp, LShrOp, LoadOp, MulOp, OrOp, PoisonOp,
44        PtrToIntOp, SDivOp, SExtOp, SIToFPOp, SRemOp, SelectOp, ShlOp, ShuffleVectorOp, StoreOp,
45        SubOp, SwitchOp, TruncOp, UDivOp, UIToFPOp, URemOp, UndefOp, XorOp, ZExtOp, ZeroOp,
46    },
47};
48
49#[derive(Error, Debug)]
50#[error("Register Promotion: Allocation info provided is not related to this operation")]
51pub struct UnrelatedAllocInfo;
52
53#[op_interface_impl]
54impl PromotableAllocationInterface for AllocaOp {
55    fn alloc_info(&self, ctx: &Context) -> Vec<AllocInfo> {
56        vec![AllocInfo {
57            ptr: self.get_result(ctx),
58            ty: self.result_pointee_type(ctx),
59        }]
60    }
61
62    fn default_value(
63        &self,
64        ctx: &mut Context,
65        inserter: &mut dyn Inserter,
66        alloc_info: &AllocInfo,
67    ) -> Result<Value> {
68        if alloc_info.ptr != self.get_result(ctx) {
69            return arg_err!(self.loc(ctx), UnrelatedAllocInfo);
70        }
71        let poison = PoisonOp::new(ctx, alloc_info.ty);
72        let poison_val = poison.get_result(ctx);
73        inserter.insert_op(ctx, &poison);
74        Ok(poison_val)
75    }
76
77    fn promote(
78        &self,
79        ctx: &mut Context,
80        rewriter: &mut dyn Rewriter,
81        alloc_infos: &[AllocInfo],
82    ) -> Result<()> {
83        if alloc_infos.len() != 1 || alloc_infos[0].ptr != self.get_result(ctx) {
84            return arg_err!(self.loc(ctx), UnrelatedAllocInfo);
85        }
86        rewriter.erase_operation(ctx, self.get_operation());
87        Ok(())
88    }
89}
90
91#[op_interface_impl]
92impl PromotableOpInterface for StoreOp {
93    fn promotion_kind(&self, ctx: &Context, alloc_info: &AllocInfo) -> PromotableOpKind {
94        if self.get_operand_address(ctx) == alloc_info.ptr {
95            PromotableOpKind::Store(self.get_operand_value(ctx))
96        } else {
97            PromotableOpKind::NonPromotableUse
98        }
99    }
100
101    fn promote(
102        &self,
103        ctx: &mut Context,
104        alloc_info_reaching_defs: &[(AllocInfo, Value)],
105        rewriter: &mut dyn Rewriter,
106    ) -> Result<()> {
107        if alloc_info_reaching_defs.len() != 1 {
108            return arg_err!(self.loc(ctx), UnrelatedAllocInfo);
109        }
110        let (alloc_info, _reaching_def) = &alloc_info_reaching_defs[0];
111        if self.get_operand_address(ctx) != alloc_info.ptr {
112            return arg_err!(self.loc(ctx), UnrelatedAllocInfo);
113        }
114        rewriter.erase_operation(ctx, self.get_operation());
115        Ok(())
116    }
117}
118
119#[op_interface_impl]
120impl PromotableOpInterface for LoadOp {
121    fn promotion_kind(&self, ctx: &Context, alloc_info: &AllocInfo) -> PromotableOpKind {
122        if self.get_operand_address(ctx) == alloc_info.ptr {
123            PromotableOpKind::Load
124        } else {
125            PromotableOpKind::NonPromotableUse
126        }
127    }
128
129    fn promote(
130        &self,
131        ctx: &mut Context,
132        alloc_info_reaching_defs: &[(AllocInfo, Value)],
133        rewriter: &mut dyn Rewriter,
134    ) -> Result<()> {
135        if alloc_info_reaching_defs.len() != 1 {
136            return arg_err!(self.loc(ctx), UnrelatedAllocInfo);
137        }
138        let (alloc_info, reaching_def) = &alloc_info_reaching_defs[0];
139        if self.get_operand_address(ctx) != alloc_info.ptr {
140            return arg_err!(self.loc(ctx), UnrelatedAllocInfo);
141        }
142        rewriter.replace_operation_with_values(ctx, self.get_operation(), vec![*reaching_def]);
143        Ok(())
144    }
145}
146
147// Implement [SideEffects] with `has_side_effects` returning `false`
148macro_rules! impl_side_effects_false {
149  ($($op:ty),+ $(,)?) => {
150    $(
151      #[op_interface_impl]
152      impl SideEffects for $op {
153        fn has_side_effects(&self, _ctx: &Context) -> bool {
154          false
155        }
156      }
157    )+
158  };
159}
160
161// Pure value-producing ops with no memory/control side effects.
162// We don't need to implement [SideEffects] for the other ops,
163// because the assumption is that the absense of the interface
164// implies the presence of side effects, which is a safe default for DCE.
165impl_side_effects_false!(
166    AddOp,
167    SubOp,
168    MulOp,
169    ShlOp,
170    UDivOp,
171    SDivOp,
172    URemOp,
173    SRemOp,
174    AndOp,
175    OrOp,
176    XorOp,
177    LShrOp,
178    AShrOp,
179    ICmpOp,
180    AllocaOp,
181    BitcastOp,
182    IntToPtrOp,
183    PtrToIntOp,
184    ConstantOp,
185    UndefOp,
186    PoisonOp,
187    FreezeOp,
188    ZeroOp,
189    AddressOfOp,
190    SExtOp,
191    ZExtOp,
192    FPExtOp,
193    TruncOp,
194    FPTruncOp,
195    FPToSIOp,
196    FPToUIOp,
197    SIToFPOp,
198    UIToFPOp,
199    InsertValueOp,
200    ExtractValueOp,
201    InsertElementOp,
202    ExtractElementOp,
203    ShuffleVectorOp,
204    SelectOp,
205    FNegOp,
206    FAddOp,
207    FSubOp,
208    FMulOp,
209    FDivOp,
210    FRemOp,
211    FCmpOp,
212    GetElementPtrOp,
213);
214
215#[op_interface_impl]
216impl BlockArgRemoval for FuncOp {
217    fn can_remove_block_args(&self, ctx: &Context, block: Ptr<BasicBlock>) -> bool {
218        !matches!(self.get_entry_block(ctx), Some(entry) if entry == block)
219    }
220}
221
222/// Assumes `operand_attrs` has length 2. If both elements are `Some(x)` where `x` can
223/// be casted to an [IntegerAttr], return the casted results. Otherwise, return `None`.
224fn get_int_bin_operands(operand_attrs: &[Option<AttrObj>]) -> Option<(IntegerAttr, IntegerAttr)> {
225    assert!(operand_attrs.len() == 2);
226    let [Some(lhs), Some(rhs)] = operand_attrs else {
227        return None;
228    };
229    let lhs_int = lhs
230        .downcast_ref::<IntegerAttr>()
231        .expect("invalid operand type: typecheck before optimizing");
232    let rhs_int = rhs
233        .downcast_ref::<IntegerAttr>()
234        .expect("invalid operand type: typecheck before optimizing");
235    Some((lhs_int.clone(), rhs_int.clone()))
236}
237
238/// Constant fold this binary integer operation, taking integer overflow flags
239/// into account.
240///
241/// `operand_attrs` contains `Some(attr)` for operands inferred constant
242/// and `None` for operands not inferred constant.
243///
244/// `flags` contains the llvm integer overflow flags associated with this operation
245///
246/// `combine` computes the wrapped result together with whether the operation
247/// unsigned- and signed-overflowed (the two booleans, in that order). The
248///
249/// Returns a singleton vector containing the folded result, or `None` if folding
250/// is not possible.
251fn check_fold_int_bin_op_with_overflow(
252    operand_attrs: &[Option<AttrObj>],
253    flags: IntegerOverflowFlagsAttr,
254    combine: impl Fn(&APInt, &APInt) -> (APInt, bool, bool),
255) -> Vec<Option<AttrObj>> {
256    let Some((lhs, rhs)) = get_int_bin_operands(operand_attrs) else {
257        return vec![None];
258    };
259    let (res, unsigned_overflow, signed_overflow) = combine(&lhs.value(), &rhs.value());
260    if (flags.nsw && signed_overflow) || (flags.nuw && unsigned_overflow) {
261        return vec![None];
262    }
263    let res = Box::new(IntegerAttr::new(lhs.get_type(), res)) as AttrObj;
264    vec![Some(res)]
265}
266
267/// Returns `true` if signed-dividing/remaindering `lhs` by `rhs` is undefined
268/// behavior in LLVM, and so must not be constant folded. The two cases are
269/// division by zero, and the signed overflow `INT_MIN / -1` (true quotient
270/// `INT_MAX + 1`, not representable), whose result LLVM leaves as poison and
271/// whose hardware behavior diverges (x86 traps, AArch64 wraps).
272fn is_signed_div_ub(lhs: &APInt, rhs: &APInt) -> bool {
273    let bw = NonZero::new(rhs.bw()).expect("operand has zero bitwidth");
274    // `-1` is the all-ones bit pattern, i.e. the unsigned max.
275    rhs.is_zero() || (*lhs == APInt::imin(bw) && *rhs == APInt::umax(bw))
276}
277
278/// Evaluate an integer comparison `lhs <pred> rhs`. `lhs` and `rhs` must have
279/// the same bitwidth.
280fn eval_icmp(pred: &ICmpPredicateAttr, lhs: &APInt, rhs: &APInt) -> bool {
281    match pred {
282        ICmpPredicateAttr::EQ => lhs == rhs,
283        ICmpPredicateAttr::NE => lhs != rhs,
284        ICmpPredicateAttr::SLT => lhs.slt(rhs),
285        ICmpPredicateAttr::SLE => lhs.sle(rhs),
286        ICmpPredicateAttr::SGT => lhs.sgt(rhs),
287        ICmpPredicateAttr::SGE => lhs.sge(rhs),
288        ICmpPredicateAttr::ULT => lhs.ult(rhs),
289        ICmpPredicateAttr::ULE => lhs.ule(rhs),
290        ICmpPredicateAttr::UGT => lhs.ugt(rhs),
291        ICmpPredicateAttr::UGE => lhs.uge(rhs),
292    }
293}
294
295/// Constant fold this binary integer operation into a singleton vector
296/// containing its result type if folding is successful, or None otherwise.
297fn check_fold_int_bin_op(
298    operand_attrs: &[Option<AttrObj>],
299    combine: impl Fn(&APInt, &APInt) -> APInt,
300) -> Vec<Option<AttrObj>> {
301    let Some((lhs, rhs)) = get_int_bin_operands(operand_attrs) else {
302        return vec![None];
303    };
304    let res = Box::new(IntegerAttr::new(
305        lhs.get_type(),
306        combine(&lhs.value(), &rhs.value()),
307    )) as AttrObj;
308    vec![Some(res)]
309}
310
311/// Returns `true` if the fast-math `flags` make a constant fold involving
312/// `values` (the operands together with the computed result) undefined
313/// behavior, so it must not be folded to a concrete value.
314fn fast_math_forbids_fold(flags: FastmathFlagsAttr, values: &[&dyn FloatAttr]) -> bool {
315    let flags = flags.0;
316    (flags.contains(FastmathFlags::NNAN) && values.iter().any(|v| v.is_nan()))
317        || (flags.contains(FastmathFlags::NINF) && values.iter().any(|v| v.is_infinite()))
318}
319
320/// Constant fold a binary floating-point operation into a singleton vector
321/// containing its result if both operands are constant, or None otherwise.
322fn check_fold_float_bin_op(
323    operand_attrs: &[Option<AttrObj>],
324    flags: FastmathFlagsAttr,
325    combine: impl Fn(&dyn FloatAttr, &dyn FloatAttr) -> Box<dyn FloatAttr>,
326) -> Vec<Option<AttrObj>> {
327    assert!(operand_attrs.len() == 2);
328    let [Some(lhs), Some(rhs)] = operand_attrs else {
329        return vec![None];
330    };
331    let lhs = attr_cast::<dyn FloatAttr>(&**lhs)
332        .expect("invalid operand type: typecheck before optimizing");
333    let rhs = attr_cast::<dyn FloatAttr>(&**rhs)
334        .expect("invalid operand type: typecheck before optimizing");
335    let res = combine(lhs, rhs);
336    if fast_math_forbids_fold(flags, &[lhs, rhs, &*res]) {
337        return vec![None];
338    }
339    let res = pliron::dyn_clone::clone_box(&*res as &dyn Attribute);
340    vec![Some(res)]
341}
342
343#[op_interface_impl]
344impl ConstFoldInterface for ConstantOp {
345    fn check_fold(
346        &self,
347        ctx: &Context,
348        _operand_attrs: &[Option<AttrObj>],
349    ) -> Vec<Option<AttrObj>> {
350        vec![Some(
351            pliron::dyn_clone::clone_box(&*self.get_value(ctx)) as AttrObj
352        )]
353    }
354
355    fn fold_in_place(
356        &self,
357        _ctx: &mut Context,
358        _operand_attrs: &[Option<AttrObj>],
359        _rewriter: &mut dyn Rewriter,
360    ) -> IRStatus {
361        IRStatus::Unchanged
362    }
363}
364
365#[op_interface_impl]
366impl ConstFoldInterface for AddOp {
367    fn check_fold(&self, ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
368        check_fold_int_bin_op_with_overflow(
369            ops,
370            self.integer_overflow_flag(ctx),
371            APInt::add_overflow,
372        )
373    }
374    fn fold_in_place(
375        &self,
376        ctx: &mut Context,
377        ops: &[Option<AttrObj>],
378        rw: &mut dyn Rewriter,
379    ) -> IRStatus {
380        self.fold_with_materialization(ctx, ops, rw)
381    }
382}
383
384#[op_interface_impl]
385impl ConstFoldInterface for SubOp {
386    fn check_fold(&self, ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
387        check_fold_int_bin_op_with_overflow(
388            ops,
389            self.integer_overflow_flag(ctx),
390            APInt::sub_overflow,
391        )
392    }
393    fn fold_in_place(
394        &self,
395        ctx: &mut Context,
396        ops: &[Option<AttrObj>],
397        rw: &mut dyn Rewriter,
398    ) -> IRStatus {
399        self.fold_with_materialization(ctx, ops, rw)
400    }
401}
402
403#[op_interface_impl]
404impl ConstFoldInterface for MulOp {
405    fn check_fold(&self, ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
406        check_fold_int_bin_op_with_overflow(
407            ops,
408            self.integer_overflow_flag(ctx),
409            APInt::mul_overflow,
410        )
411    }
412    fn fold_in_place(
413        &self,
414        ctx: &mut Context,
415        ops: &[Option<AttrObj>],
416        rw: &mut dyn Rewriter,
417    ) -> IRStatus {
418        self.fold_with_materialization(ctx, ops, rw)
419    }
420}
421
422#[op_interface_impl]
423impl ConstFoldInterface for ShlOp {
424    fn check_fold(&self, ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
425        match get_int_bin_operands(ops) {
426            Some((lhs, rhs)) => {
427                let shamt = rhs.value();
428                let lhs_bw: usize = lhs.value().bw();
429                let lhs_bw: APInt = APInt::from_usize(lhs_bw, NonZero::new(lhs_bw).unwrap());
430                if shamt.ult(&lhs_bw) {
431                    check_fold_int_bin_op_with_overflow(
432                        ops,
433                        self.integer_overflow_flag(ctx),
434                        APInt::shl_overflow,
435                    )
436                } else {
437                    vec![None]
438                }
439            }
440            None => vec![None],
441        }
442    }
443    fn fold_in_place(
444        &self,
445        ctx: &mut Context,
446        ops: &[Option<AttrObj>],
447        rw: &mut dyn Rewriter,
448    ) -> IRStatus {
449        self.fold_with_materialization(ctx, ops, rw)
450    }
451}
452
453#[op_interface_impl]
454impl ConstFoldInterface for UDivOp {
455    fn check_fold(&self, _ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
456        match get_int_bin_operands(ops) {
457            Some((_, rhs)) if rhs.value().is_zero() => vec![None],
458            _ => check_fold_int_bin_op(ops, APInt::udiv),
459        }
460    }
461    fn fold_in_place(
462        &self,
463        ctx: &mut Context,
464        ops: &[Option<AttrObj>],
465        rw: &mut dyn Rewriter,
466    ) -> IRStatus {
467        self.fold_with_materialization(ctx, ops, rw)
468    }
469}
470
471#[op_interface_impl]
472impl ConstFoldInterface for SDivOp {
473    fn check_fold(&self, _ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
474        match get_int_bin_operands(ops) {
475            Some((lhs, rhs)) if is_signed_div_ub(&lhs.value(), &rhs.value()) => vec![None],
476            _ => check_fold_int_bin_op(ops, APInt::sdiv),
477        }
478    }
479    fn fold_in_place(
480        &self,
481        ctx: &mut Context,
482        ops: &[Option<AttrObj>],
483        rw: &mut dyn Rewriter,
484    ) -> IRStatus {
485        self.fold_with_materialization(ctx, ops, rw)
486    }
487}
488
489#[op_interface_impl]
490impl ConstFoldInterface for URemOp {
491    fn check_fold(&self, _ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
492        match get_int_bin_operands(ops) {
493            Some((_, rhs)) if rhs.value().is_zero() => vec![None],
494            _ => check_fold_int_bin_op(ops, APInt::urem),
495        }
496    }
497    fn fold_in_place(
498        &self,
499        ctx: &mut Context,
500        ops: &[Option<AttrObj>],
501        rw: &mut dyn Rewriter,
502    ) -> IRStatus {
503        self.fold_with_materialization(ctx, ops, rw)
504    }
505}
506
507#[op_interface_impl]
508impl ConstFoldInterface for SRemOp {
509    fn check_fold(&self, _ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
510        match get_int_bin_operands(ops) {
511            Some((lhs, rhs)) if is_signed_div_ub(&lhs.value(), &rhs.value()) => vec![None],
512            _ => check_fold_int_bin_op(ops, APInt::srem),
513        }
514    }
515    fn fold_in_place(
516        &self,
517        ctx: &mut Context,
518        ops: &[Option<AttrObj>],
519        rw: &mut dyn Rewriter,
520    ) -> IRStatus {
521        self.fold_with_materialization(ctx, ops, rw)
522    }
523}
524
525#[op_interface_impl]
526impl ConstFoldInterface for AndOp {
527    fn check_fold(&self, _ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
528        assert!(ops.len() == 2);
529        for op in ops.iter().flatten() {
530            let int = op
531                .downcast_ref::<IntegerAttr>()
532                .expect("invalid operand type: typecheck before optimizing");
533            if int.value().is_zero() {
534                let zero = APInt::zero(NonZero::new(int.value().bw()).expect("zero bitwidth"));
535                let res = Box::new(IntegerAttr::new(int.get_type(), zero)) as AttrObj;
536                return vec![Some(res)];
537            }
538        }
539        check_fold_int_bin_op(ops, APInt::and)
540    }
541
542    fn fold_in_place(
543        &self,
544        ctx: &mut Context,
545        ops: &[Option<AttrObj>],
546        rw: &mut dyn Rewriter,
547    ) -> IRStatus {
548        self.fold_with_materialization(ctx, ops, rw)
549    }
550}
551
552#[op_interface_impl]
553impl ConstFoldInterface for OrOp {
554    fn check_fold(&self, _ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
555        assert!(ops.len() == 2);
556        for op in ops.iter().flatten() {
557            let int = op
558                .downcast_ref::<IntegerAttr>()
559                .expect("invalid operand type: typecheck before optimizing");
560            let bw = NonZero::new(int.value().bw()).expect("zero bitwidth");
561            if int.value() == APInt::umax(bw) {
562                let all_ones = APInt::umax(bw);
563                let res = Box::new(IntegerAttr::new(int.get_type(), all_ones)) as AttrObj;
564                return vec![Some(res)];
565            }
566        }
567        check_fold_int_bin_op(ops, APInt::or)
568    }
569
570    fn fold_in_place(
571        &self,
572        ctx: &mut Context,
573        ops: &[Option<AttrObj>],
574        rw: &mut dyn Rewriter,
575    ) -> IRStatus {
576        self.fold_with_materialization(ctx, ops, rw)
577    }
578}
579
580#[op_interface_impl]
581impl ConstFoldInterface for XorOp {
582    fn check_fold(&self, _ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
583        check_fold_int_bin_op(ops, APInt::xor)
584    }
585
586    fn fold_in_place(
587        &self,
588        ctx: &mut Context,
589        ops: &[Option<AttrObj>],
590        rw: &mut dyn Rewriter,
591    ) -> IRStatus {
592        self.fold_with_materialization(ctx, ops, rw)
593    }
594}
595
596#[op_interface_impl]
597impl ConstFoldInterface for LShrOp {
598    fn check_fold(&self, _ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
599        // A shift amount >= the bitwidth is undefined behavior in LLVM, so it
600        // must not be folded.
601        match get_int_bin_operands(ops) {
602            Some((lhs, rhs)) => {
603                let lhs_bw = lhs.value().bw();
604                let lhs_bw = APInt::from_usize(lhs_bw, NonZero::new(lhs_bw).unwrap());
605                if rhs.value().ult(&lhs_bw) {
606                    check_fold_int_bin_op(ops, APInt::lshr)
607                } else {
608                    vec![None]
609                }
610            }
611            None => vec![None],
612        }
613    }
614    fn fold_in_place(
615        &self,
616        ctx: &mut Context,
617        ops: &[Option<AttrObj>],
618        rw: &mut dyn Rewriter,
619    ) -> IRStatus {
620        self.fold_with_materialization(ctx, ops, rw)
621    }
622}
623
624#[op_interface_impl]
625impl ConstFoldInterface for AShrOp {
626    fn check_fold(&self, _ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
627        // A shift amount >= the bitwidth is undefined behavior in LLVM, so it
628        // must not be folded.
629        match get_int_bin_operands(ops) {
630            Some((lhs, rhs)) => {
631                let lhs_bw = lhs.value().bw();
632                let lhs_bw = APInt::from_usize(lhs_bw, NonZero::new(lhs_bw).unwrap());
633                if rhs.value().ult(&lhs_bw) {
634                    check_fold_int_bin_op(ops, APInt::ashr)
635                } else {
636                    vec![None]
637                }
638            }
639            None => vec![None],
640        }
641    }
642    fn fold_in_place(
643        &self,
644        ctx: &mut Context,
645        ops: &[Option<AttrObj>],
646        rw: &mut dyn Rewriter,
647    ) -> IRStatus {
648        self.fold_with_materialization(ctx, ops, rw)
649    }
650}
651
652#[op_interface_impl]
653impl ConstFoldInterface for ICmpOp {
654    fn check_fold(&self, ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
655        let Some((lhs, rhs)) = get_int_bin_operands(ops) else {
656            return vec![None];
657        };
658        let result = eval_icmp(&self.predicate(ctx), &lhs.value(), &rhs.value());
659        let bool_ty = IntegerType::get(ctx, 1, Signedness::Signless);
660        let res = Box::new(IntegerAttr::new(
661            bool_ty,
662            APInt::from_u8(result as u8, bw(1)),
663        )) as AttrObj;
664        vec![Some(res)]
665    }
666    fn fold_in_place(
667        &self,
668        ctx: &mut Context,
669        ops: &[Option<AttrObj>],
670        rw: &mut dyn Rewriter,
671    ) -> IRStatus {
672        self.fold_with_materialization(ctx, ops, rw)
673    }
674}
675
676#[op_interface_impl]
677impl ConstFoldInterface for SExtOp {
678    fn check_fold(&self, ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
679        let [Some(operand)] = ops else {
680            return vec![None];
681        };
682        let operand = operand
683            .downcast_ref::<IntegerAttr>()
684            .expect("invalid operand type: typecheck before optimizing");
685        let res_ty = self.result_type(ctx);
686        let dest_width = res_ty
687            .deref(ctx)
688            .downcast_ref::<IntegerType>()
689            .expect("sext result must be an integer type")
690            .width();
691        let dest_ty = IntegerType::get(ctx, dest_width, Signedness::Signless);
692        let extended = operand
693            .value()
694            .sext(NonZero::new(dest_width as usize).expect("result has zero bitwidth"));
695        let res = Box::new(IntegerAttr::new(dest_ty, extended)) as AttrObj;
696        vec![Some(res)]
697    }
698    fn fold_in_place(
699        &self,
700        ctx: &mut Context,
701        ops: &[Option<AttrObj>],
702        rw: &mut dyn Rewriter,
703    ) -> IRStatus {
704        self.fold_with_materialization(ctx, ops, rw)
705    }
706}
707
708#[op_interface_impl]
709impl ConstFoldInterface for ZExtOp {
710    fn check_fold(&self, ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
711        let [Some(operand)] = ops else {
712            return vec![None];
713        };
714        let operand = operand
715            .downcast_ref::<IntegerAttr>()
716            .expect("invalid operand type: typecheck before optimizing");
717        // `zext nneg` asserts the operand is non-negative; if it isn't, the
718        // result is poison, so we must not fold it to a concrete value.
719        let value = operand.value();
720        if self.nneg(ctx)
721            && value.slt(&APInt::zero(
722                NonZero::new(value.bw()).expect("operand has zero bitwidth"),
723            ))
724        {
725            return vec![None];
726        }
727        let res_ty = self.result_type(ctx);
728        let dest_width = res_ty
729            .deref(ctx)
730            .downcast_ref::<IntegerType>()
731            .expect("zext result must be an integer type")
732            .width();
733        let dest_ty = IntegerType::get(ctx, dest_width, Signedness::Signless);
734        let extended =
735            value.zext(NonZero::new(dest_width as usize).expect("result has zero bitwidth"));
736        let res = Box::new(IntegerAttr::new(dest_ty, extended)) as AttrObj;
737        vec![Some(res)]
738    }
739    fn fold_in_place(
740        &self,
741        ctx: &mut Context,
742        ops: &[Option<AttrObj>],
743        rw: &mut dyn Rewriter,
744    ) -> IRStatus {
745        self.fold_with_materialization(ctx, ops, rw)
746    }
747}
748
749#[op_interface_impl]
750impl ConstFoldInterface for FNegOp {
751    fn check_fold(&self, ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
752        let [Some(operand)] = ops else {
753            return vec![None];
754        };
755        let float_val = attr_cast::<dyn FloatAttr>(&**operand)
756            .expect("invalid operand type: typecheck before optimizing");
757        let negated = float_val.neg();
758        // Negation cannot create or destroy NaN/Inf, so checking the operand
759        // covers the result too.
760        if fast_math_forbids_fold(self.fast_math_flags(ctx), &[float_val]) {
761            return vec![None];
762        }
763        let res = pliron::dyn_clone::clone_box(&*negated as &dyn Attribute);
764        vec![Some(res)]
765    }
766    fn fold_in_place(
767        &self,
768        ctx: &mut Context,
769        ops: &[Option<AttrObj>],
770        rw: &mut dyn Rewriter,
771    ) -> IRStatus {
772        self.fold_with_materialization(ctx, ops, rw)
773    }
774}
775
776#[op_interface_impl]
777impl ConstFoldInterface for FAddOp {
778    fn check_fold(&self, ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
779        check_fold_float_bin_op(ops, self.fast_math_flags(ctx), |lhs, rhs| {
780            lhs.add(rhs).value
781        })
782    }
783    fn fold_in_place(
784        &self,
785        ctx: &mut Context,
786        ops: &[Option<AttrObj>],
787        rw: &mut dyn Rewriter,
788    ) -> IRStatus {
789        self.fold_with_materialization(ctx, ops, rw)
790    }
791}
792
793#[op_interface_impl]
794impl ConstFoldInterface for FSubOp {
795    fn check_fold(&self, ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
796        check_fold_float_bin_op(ops, self.fast_math_flags(ctx), |lhs, rhs| {
797            lhs.sub(rhs).value
798        })
799    }
800    fn fold_in_place(
801        &self,
802        ctx: &mut Context,
803        ops: &[Option<AttrObj>],
804        rw: &mut dyn Rewriter,
805    ) -> IRStatus {
806        self.fold_with_materialization(ctx, ops, rw)
807    }
808}
809
810#[op_interface_impl]
811impl ConstFoldInterface for FMulOp {
812    fn check_fold(&self, ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
813        check_fold_float_bin_op(ops, self.fast_math_flags(ctx), |lhs, rhs| {
814            lhs.mul(rhs).value
815        })
816    }
817    fn fold_in_place(
818        &self,
819        ctx: &mut Context,
820        ops: &[Option<AttrObj>],
821        rw: &mut dyn Rewriter,
822    ) -> IRStatus {
823        self.fold_with_materialization(ctx, ops, rw)
824    }
825}
826
827#[op_interface_impl]
828impl ConstFoldInterface for FDivOp {
829    fn check_fold(&self, ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
830        check_fold_float_bin_op(ops, self.fast_math_flags(ctx), |lhs, rhs| {
831            lhs.div(rhs).value
832        })
833    }
834    fn fold_in_place(
835        &self,
836        ctx: &mut Context,
837        ops: &[Option<AttrObj>],
838        rw: &mut dyn Rewriter,
839    ) -> IRStatus {
840        self.fold_with_materialization(ctx, ops, rw)
841    }
842}
843
844#[op_interface_impl]
845impl ConstFoldInterface for FRemOp {
846    fn check_fold(&self, ctx: &Context, ops: &[Option<AttrObj>]) -> Vec<Option<AttrObj>> {
847        check_fold_float_bin_op(ops, self.fast_math_flags(ctx), |lhs, rhs| {
848            lhs.rem(rhs).value
849        })
850    }
851    fn fold_in_place(
852        &self,
853        ctx: &mut Context,
854        ops: &[Option<AttrObj>],
855        rw: &mut dyn Rewriter,
856    ) -> IRStatus {
857        self.fold_with_materialization(ctx, ops, rw)
858    }
859}
860
861#[op_interface_impl]
862impl BranchOpFoldInterface for BrOp {
863    fn check_fold(&self, ctx: &Context, _operands: &[Option<AttrObj>]) -> Vec<Ptr<BasicBlock>> {
864        self.get_operation().deref(ctx).successors().collect()
865    }
866    fn fold_in_place(
867        &self,
868        _ctx: &mut Context,
869        _ops: &[Option<AttrObj>],
870        _rw: &mut dyn Rewriter,
871    ) -> IRStatus {
872        IRStatus::Unchanged
873    }
874}
875
876impl CondBrOp {
877    fn possible_successor_indices(
878        &self,
879        ctx: &Context,
880        operands: &[Option<AttrObj>],
881    ) -> Vec<usize> {
882        let Some(cond_attr) = operands.first().unwrap().as_ref() else {
883            let num_successors = self.get_operation().deref(ctx).successors().count();
884            return (0..num_successors).collect();
885        };
886        let cond_int = cond_attr
887            .downcast_ref::<IntegerAttr>()
888            .expect("CondBrOp condition operand must be an IntegerAttr");
889        let taken = if cond_int.value().is_zero() { 1 } else { 0 };
890        vec![taken]
891    }
892}
893
894#[op_interface_impl]
895impl BranchOpFoldInterface for CondBrOp {
896    fn check_fold(&self, ctx: &Context, operands: &[Option<AttrObj>]) -> Vec<Ptr<BasicBlock>> {
897        let successors: Vec<Ptr<BasicBlock>> =
898            self.get_operation().deref(ctx).successors().collect();
899
900        self.possible_successor_indices(ctx, operands)
901            .iter()
902            .map(|ind| successors[*ind])
903            .collect()
904    }
905
906    fn fold_in_place(
907        &self,
908        ctx: &mut Context,
909        ops: &[Option<AttrObj>],
910        rewriter: &mut dyn Rewriter,
911    ) -> IRStatus {
912        let possible_successor_indices = self.possible_successor_indices(ctx, ops);
913        if possible_successor_indices.len() != 1 {
914            return IRStatus::Unchanged;
915        };
916        let successor_ind = possible_successor_indices[0];
917        let successors: Vec<Ptr<BasicBlock>> =
918            self.get_operation().deref(ctx).successors().collect();
919        let new_op = BrOp::new(
920            ctx,
921            successors[successor_ind],
922            self.successor_operands(ctx, successor_ind),
923        )
924        .get_operation();
925        let old_op = self.get_operation();
926        rewriter.insert_operation(ctx, new_op);
927        rewriter.replace_operation(ctx, old_op, new_op);
928        IRStatus::Changed
929    }
930}
931
932#[op_interface_impl]
933impl BranchOpFoldInterface for SwitchOp {
934    fn check_fold(&self, ctx: &Context, operands: &[Option<AttrObj>]) -> Vec<Ptr<BasicBlock>> {
935        let successors: Vec<Ptr<BasicBlock>> =
936            self.get_operation().deref(ctx).successors().collect();
937        let Some(cond_attr) = operands.first().and_then(|o| o.as_ref()) else {
938            return successors;
939        };
940        let cond_int = cond_attr
941            .downcast_ref::<IntegerAttr>()
942            .expect("Switch condition operand must be an IntegerAttr")
943            .value();
944        // Successor 0 is the default destination; successors 1..N correspond to case_values[0..N-1].
945        let case_values = self
946            .get_attr_llvm_switch_case_values(ctx)
947            .expect("SwitchOp missing case values attribute");
948        let taken = case_values
949            .0
950            .iter()
951            .position(|case| case.value() == cond_int)
952            .map(|i| i + 1)
953            .unwrap_or(0);
954        vec![successors[taken]]
955    }
956
957    fn fold_in_place(
958        &self,
959        ctx: &mut Context,
960        ops: &[Option<AttrObj>],
961        rewriter: &mut dyn Rewriter,
962    ) -> IRStatus {
963        let Some(cond_attr) = ops.first().unwrap().as_ref() else {
964            return IRStatus::Unchanged;
965        };
966        let cond_int = cond_attr
967            .downcast_ref::<IntegerAttr>()
968            .expect("Switch condition operand must be an IntegerAttr")
969            .value();
970        let successor_ind = {
971            let case_values = self
972                .get_attr_llvm_switch_case_values(ctx)
973                .expect("SwitchOp missing case values attribute");
974            case_values
975                .0
976                .iter()
977                .position(|case| case.value() == cond_int)
978                // There is no case value corresponding to the default successor,
979                // so case_values index 0 corresponds to succesors index 1, etc.
980                .map(|i| i + 1)
981                // successor index 0 is the default successor
982                .unwrap_or(0)
983        };
984        let successors: Vec<Ptr<BasicBlock>> =
985            self.get_operation().deref(ctx).successors().collect();
986        let new_op = BrOp::new(
987            ctx,
988            successors[successor_ind],
989            self.successor_operands(ctx, successor_ind),
990        )
991        .get_operation();
992        let old_op = self.get_operation();
993        rewriter.insert_operation(ctx, new_op);
994        rewriter.replace_operation(ctx, old_op, new_op);
995        IRStatus::Changed
996    }
997}
pliron_llvm/interface_impls.rs

pliron_llvm/
interface_impls.rs
