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

pliron_llvm/
interface_impls.rs
