Skip to main content

pliron/opts/constants/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) The pliron contributors
3
4use alloc::vec::Vec;
5use pliron_derive::op_interface;
6
7use crate::{
8    attribute::{AttrObj, attr_cast},
9    basic_block::BasicBlock,
10    builtin::{attr_interfaces::MaterializableAttr, op_interfaces::BranchOpInterface},
11    context::{Context, Ptr},
12    irbuild::{IRStatus, rewriter::Rewriter},
13    op::Op,
14    result::Result,
15};
16
17pub mod sccp;
18mod state;
19
20/// Interface for constant folding of operations.
21#[op_interface]
22pub trait ConstFoldInterface {
23    /// Given a slice `operand_attrs` corresponding to each operand, indicating a known
24    /// compile time constant value for that operand (if any), returns a vector corresponding
25    /// to each result, indicating the folded (inferred constant) value, if any.
26    fn check_fold(&self, ctx: &Context, operand_attrs: &[Option<AttrObj>]) -> Vec<Option<AttrObj>>;
27
28    /// Given a slice `operand_attrs` corresponding to each operand, indicating a known
29    /// compile time constant value for that operand (if any), attempts to fold the op in
30    /// place using the provided `rewriter`. Assumes that `rewriter` is positioned just
31    /// before the op to be folded.
32    ///
33    /// Implementors that fold by materializing constant values for their results can
34    /// usually delegate to [fold_with_materialization](Self::fold_with_materialization)
35    /// rather than reimplementing the materialization logic.
36    fn fold_in_place(
37        &self,
38        ctx: &mut Context,
39        operand_attrs: &[Option<AttrObj>],
40        rewriter: &mut dyn Rewriter,
41    ) -> IRStatus;
42
43    /// A helper for implementing [fold_in_place](Self::fold_in_place) by materializing
44    /// the constants inferred by [check_fold](Self::check_fold).
45    ///
46    /// Only constants whose attribute types implement [MaterializableAttr] get
47    /// materialized.
48    fn fold_with_materialization(
49        &self,
50        ctx: &mut Context,
51        operand_attrs: &[Option<AttrObj>],
52        rewriter: &mut dyn Rewriter,
53    ) -> IRStatus {
54        let folded = self.check_fold(ctx, operand_attrs);
55        let op = self.get_operation();
56
57        let mut status = IRStatus::Unchanged;
58        for (result_idx, attr) in folded.iter().enumerate() {
59            let Some(attr) = attr else {
60                continue;
61            };
62            let Some(materializable) = attr_cast::<dyn MaterializableAttr>(&**attr) else {
63                log::info!(
64                    "Constant propagation tried to materialize {}, but its type does not \
65                     implement MaterializableAttr. This potentially prevents optimizations.",
66                    attr.disp(ctx)
67                );
68                continue;
69            };
70            let const_op = materializable.materialize(ctx);
71            rewriter.append_operation(ctx, const_op);
72            let new_value = const_op.deref(ctx).get_result(0);
73            let old_value = op.deref(ctx).get_result(result_idx);
74            rewriter.replace_value_uses_with(ctx, old_value, new_value);
75            status = IRStatus::Changed;
76        }
77        status
78    }
79
80    fn verify(_op: &dyn Op, _ctx: &Context) -> Result<()>
81    where
82        Self: Sized,
83    {
84        Ok(())
85    }
86}
87
88/// Interface for ruling out branch destinations
89/// based on static information about branch conditions.
90#[op_interface]
91pub trait BranchOpFoldInterface: BranchOpInterface {
92    /// Return the list of possible successor blocks given that `operands`
93    /// contains `Some(attr)` for each operand known to be constant, where `attr` contains
94    /// the known constant value.
95    fn check_fold(&self, ctx: &Context, operands: &[Option<AttrObj>]) -> Vec<Ptr<BasicBlock>>;
96
97    /// Given a slice `operand_attrs` corresponding to each operand, indicating a known
98    /// compile time constant value for that operand (if any), attempts to fold the op in
99    /// place using the provided `rewriter`. Assumes that `rewriter` is positioned just
100    /// before the op to be folded.
101    fn fold_in_place(
102        &self,
103        ctx: &mut Context,
104        operand_attrs: &[Option<AttrObj>],
105        rewriter: &mut dyn Rewriter,
106    ) -> IRStatus;
107
108    fn verify(_op: &dyn Op, _ctx: &Context) -> Result<()>
109    where
110        Self: Sized,
111    {
112        Ok(())
113    }
114}