feat: constant folding implemented for core and float extension (#769)

Closes #711 

Original review in #758

---------

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Alan Lawrence <[email protected]>
Co-authored-by: Luca Mondada <[email protected]>
Co-authored-by: Alec Edgington <[email protected]>
Co-authored-by: Alan Lawrence <[email protected]>
Co-authored-by: Agustín Borgna <[email protected]>
Co-authored-by: Luca Mondada <[email protected]>

src/algorithm.rs

@@ -1,4 +1,5 @@
 //! Algorithms using the Hugr.
 
+pub mod const_fold;
 mod half_node;
 pub mod nest_cfgs;

src/algorithm/const_fold.rs

@@ -0,0 +1,303 @@
+//! Constant folding routines.
+
+use std::collections::{BTreeSet, HashMap};
+
+use itertools::Itertools;
+
+use crate::{
+    builder::{DFGBuilder, Dataflow, DataflowHugr},
+    extension::{ConstFoldResult, ExtensionRegistry},
+    hugr::{
+        rewrite::consts::{RemoveConst, RemoveConstIgnore},
+        views::SiblingSubgraph,
+        HugrMut,
+    },
+    ops::{Const, LeafOp, OpType},
+    type_row,
+    types::{FunctionType, Type, TypeEnum},
+    values::Value,
+    Hugr, HugrView, IncomingPort, Node, SimpleReplacement,
+};
+
+/// Tag some output constants with [`OutgoingPort`] inferred from the ordering.
+fn out_row(consts: impl IntoIterator<Item = Const>) -> ConstFoldResult {
+    let vec = consts
+        .into_iter()
+        .enumerate()
+        .map(|(i, c)| (i.into(), c))
+        .collect();
+    Some(vec)
+}
+
+/// Sort folding inputs with [`IncomingPort`] as key
+fn sort_by_in_port(consts: &[(IncomingPort, Const)]) -> Vec<&(IncomingPort, Const)> {
+    let mut v: Vec<_> = consts.iter().collect();
+    v.sort_by_key(|(i, _)| i);
+    v
+}
+
+/// Sort some input constants by port and just return the constants.
+pub(crate) fn sorted_consts(consts: &[(IncomingPort, Const)]) -> Vec<&Const> {
+    sort_by_in_port(consts)
+        .into_iter()
+        .map(|(_, c)| c)
+        .collect()
+}
+/// For a given op and consts, attempt to evaluate the op.
+pub fn fold_const(op: &OpType, consts: &[(IncomingPort, Const)]) -> ConstFoldResult {
+    let op = op.as_leaf_op()?;
+
+    match op {
+        LeafOp::Noop { .. } => out_row([consts.first()?.1.clone()]),
+        LeafOp::MakeTuple { .. } => {
+            out_row([Const::new_tuple(sorted_consts(consts).into_iter().cloned())])
+        }
+        LeafOp::UnpackTuple { .. } => {
+            let c = &consts.first()?.1;
+
+            if let Value::Tuple { vs } = c.value() {
+                if let TypeEnum::Tuple(tys) = c.const_type().as_type_enum() {
+                    return out_row(tys.iter().zip(vs.iter()).map(|(t, v)| {
+                        Const::new(v.clone(), t.clone())
+                            .expect("types should already have been checked")
+                    }));
+                }
+            }
+            panic!("This op always takes a Tuple input.");
+        }
+
+        LeafOp::Tag { tag, variants } => out_row([Const::new(
+            Value::sum(*tag, consts.first()?.1.value().clone()),
+            Type::new_sum(variants.clone()),
+        )
+        .unwrap()]),
+        LeafOp::CustomOp(_) => {
+            let ext_op = op.as_extension_op()?;
+
+            ext_op.constant_fold(consts)
+        }
+        _ => None,
+    }
+}
+
+/// Generate a graph that loads and outputs `consts` in order, validating
+/// against `reg`.
+fn const_graph(consts: Vec<Const>, reg: &ExtensionRegistry) -> Hugr {
+    let const_types = consts.iter().map(Const::const_type).cloned().collect_vec();
+    let mut b = DFGBuilder::new(FunctionType::new(type_row![], const_types)).unwrap();
+
+    let outputs = consts
+        .into_iter()
+        .map(|c| b.add_load_const(c).unwrap())
+        .collect_vec();
+
+    b.finish_hugr_with_outputs(outputs, reg).unwrap()
+}
+
+/// Given some `candidate_nodes` to search for LoadConstant operations in `hugr`,
+/// return an iterator of possible constant folding rewrites. The
+/// [`SimpleReplacement`] replaces an operation with constants that result from
+/// evaluating it, the extension registry `reg` is used to validate the
+/// replacement HUGR. The vector of [`RemoveConstIgnore`] refer to the
+/// LoadConstant nodes that could be removed - they are not automatically
+/// removed as they may be used by other operations.
+pub fn find_consts<'a, 'r: 'a>(
+    hugr: &'a impl HugrView,
+    candidate_nodes: impl IntoIterator<Item = Node> + 'a,
+    reg: &'r ExtensionRegistry,
+) -> impl Iterator<Item = (SimpleReplacement, Vec<RemoveConstIgnore>)> + 'a {
+    // track nodes for operations that have already been considered for folding
+    let mut used_neighbours = BTreeSet::new();
+
+    candidate_nodes
+        .into_iter()
+        .filter_map(move |n| {
+            // only look at LoadConstant
+            hugr.get_optype(n).is_load_constant().then_some(())?;
+
+            let (out_p, _) = hugr.out_value_types(n).exactly_one().ok()?;
+            let neighbours = hugr
+                .linked_inputs(n, out_p)
+                .filter(|(n, _)| used_neighbours.insert(*n))
+                .collect_vec();
+            if neighbours.is_empty() {
+                // no uses of LoadConstant that haven't already been considered.
+                return None;
+            }
+            let fold_iter = neighbours
+                .into_iter()
+                .filter_map(|(neighbour, _)| fold_op(hugr, neighbour, reg));
+            Some(fold_iter)
+        })
+        .flatten()
+}
+
+/// Attempt to evaluate and generate rewrites for the operation at `op_node`
+fn fold_op(
+    hugr: &impl HugrView,
+    op_node: Node,
+    reg: &ExtensionRegistry,
+) -> Option<(SimpleReplacement, Vec<RemoveConstIgnore>)> {
+    let (in_consts, removals): (Vec<_>, Vec<_>) = hugr
+        .node_inputs(op_node)
+        .filter_map(|in_p| {
+            let (con_op, load_n) = get_const(hugr, op_node, in_p)?;
+            Some(((in_p, con_op), RemoveConstIgnore(load_n)))
+        })
+        .unzip();
+    let neighbour_op = hugr.get_optype(op_node);
+    // attempt to evaluate op
+    let folded = fold_const(neighbour_op, &in_consts)?;
+    let (op_outs, consts): (Vec<_>, Vec<_>) = folded.into_iter().unzip();
+    let nu_out = op_outs
+        .into_iter()
+        .enumerate()
+        .filter_map(|(i, out)| {
+            // map from the ports the op was linked to, to the output ports of
+            // the replacement.
+            hugr.single_linked_input(op_node, out)
+                .map(|np| (np, i.into()))
+        })
+        .collect();
+    let replacement = const_graph(consts, reg);
+    let sibling_graph = SiblingSubgraph::try_from_nodes([op_node], hugr)
+        .expect("Operation should form valid subgraph.");
+
+    let simple_replace = SimpleReplacement::new(
+        sibling_graph,
+        replacement,
+        // no inputs to replacement
+        HashMap::new(),
+        nu_out,
+    );
+    Some((simple_replace, removals))
+}
+
+/// If `op_node` is connected to a LoadConstant at `in_p`, return the constant
+/// and the LoadConstant node
+fn get_const(hugr: &impl HugrView, op_node: Node, in_p: IncomingPort) -> Option<(Const, Node)> {
+    let (load_n, _) = hugr.single_linked_output(op_node, in_p)?;
+    let load_op = hugr.get_optype(load_n).as_load_constant()?;
+    let const_node = hugr
+        .linked_outputs(load_n, load_op.constant_port())
+        .exactly_one()
+        .ok()?
+        .0;
+
+    let const_op = hugr.get_optype(const_node).as_const()?;
+
+    // TODO avoid const clone here
+    Some((const_op.clone(), load_n))
+}
+
+/// Exhaustively apply constant folding to a HUGR.
+pub fn constant_fold_pass(h: &mut impl HugrMut, reg: &ExtensionRegistry) {
+    loop {
+        // would be preferable if the candidates were updated to be just the
+        // neighbouring nodes of those added.
+        let rewrites = find_consts(h, h.nodes(), reg).collect_vec();
+        if rewrites.is_empty() {
+            break;
+        }
+        for (replace, removes) in rewrites {
+            h.apply_rewrite(replace).unwrap();
+            for rem in removes {
+                if let Ok(const_node) = h.apply_rewrite(rem) {
+                    // if the LoadConst was removed, try removing the Const too.
+                    if h.apply_rewrite(RemoveConst(const_node)).is_err() {
+                        // const cannot be removed - no problem
+                        continue;
+                    }
+                }
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+
+    use crate::extension::{ExtensionRegistry, PRELUDE};
+    use crate::std_extensions::arithmetic;
+
+    use crate::std_extensions::arithmetic::float_ops::FloatOps;
+    use crate::std_extensions::arithmetic::float_types::{ConstF64, FLOAT64_TYPE};
+
+    use rstest::rstest;
+
+    use super::*;
+
+    /// float to constant
+    fn f2c(f: f64) -> Const {
+        ConstF64::new(f).into()
+    }
+
+    #[rstest]
+    #[case(0.0, 0.0, 0.0)]
+    #[case(0.0, 1.0, 1.0)]
+    #[case(23.5, 435.5, 459.0)]
+    // c = a + b
+    fn test_add(#[case] a: f64, #[case] b: f64, #[case] c: f64) {
+        let consts = vec![(0.into(), f2c(a)), (1.into(), f2c(b))];
+        let add_op: OpType = FloatOps::fadd.into();
+        let out = fold_const(&add_op, &consts).unwrap();
+
+        assert_eq!(&out[..], &[(0.into(), f2c(c))]);
+    }
+
+    #[test]
+    fn test_big() {
+        /*
+           Test hugr approximately calculates
+           let x = (5.5, 3.25);
+           x.0 - x.1 == 2.25
+        */
+        let mut build =
+            DFGBuilder::new(FunctionType::new(type_row![], type_row![FLOAT64_TYPE])).unwrap();
+
+        let tup = build
+            .add_load_const(Const::new_tuple([f2c(5.5), f2c(3.25)]))
+            .unwrap();
+
+        let unpack = build
+            .add_dataflow_op(
+                LeafOp::UnpackTuple {
+                    tys: type_row![FLOAT64_TYPE, FLOAT64_TYPE],
+                },
+                [tup],
+            )
+            .unwrap();
+
+        let sub = build
+            .add_dataflow_op(FloatOps::fsub, unpack.outputs())
+            .unwrap();
+
+        let reg = ExtensionRegistry::try_new([
+            PRELUDE.to_owned(),
+            arithmetic::float_types::EXTENSION.to_owned(),
+            arithmetic::float_ops::EXTENSION.to_owned(),
+        ])
+        .unwrap();
+        let mut h = build.finish_hugr_with_outputs(sub.outputs(), &reg).unwrap();
+        assert_eq!(h.node_count(), 7);
+
+        constant_fold_pass(&mut h, &reg);
+
+        assert_fully_folded(&h, &f2c(2.25));
+    }
+    fn assert_fully_folded(h: &Hugr, expected_const: &Const) {
+        // check the hugr just loads and returns a single const
+        let mut node_count = 0;
+
+        for node in h.children(h.root()) {
+            let op = h.get_optype(node);
+            match op {
+                OpType::Input(_) | OpType::Output(_) | OpType::LoadConstant(_) => node_count += 1,
+                OpType::Const(c) if c == expected_const => node_count += 1,
+                _ => panic!("unexpected op: {:?}", op),
+            }
+        }
+
+        assert_eq!(node_count, 4);
+    }
+}

src/extension.rs

@@ -28,9 +28,11 @@ pub use op_def::{
 };
 mod type_def;
 pub use type_def::{TypeDef, TypeDefBound};
+mod const_fold;
 pub mod prelude;
 pub mod simple_op;
 pub mod validate;
+pub use const_fold::{ConstFold, ConstFoldResult};
 pub use prelude::{PRELUDE, PRELUDE_REGISTRY};
 
 /// Extension Registries store extensions to be looked up e.g. during validation.

src/extension/const_fold.rs

@@ -0,0 +1,53 @@
+use std::fmt::Formatter;
+
+use std::fmt::Debug;
+
+use crate::types::TypeArg;
+
+use crate::OutgoingPort;
+
+use crate::ops;
+
+/// Output of constant folding an operation, None indicates folding was either
+/// not possible or unsuccessful. An empty vector indicates folding was
+/// successful and no values are output.
+pub type ConstFoldResult = Option<Vec<(OutgoingPort, ops::Const)>>;
+
+/// Trait implemented by extension operations that can perform constant folding.
+pub trait ConstFold: Send + Sync {
+    /// Given type arguments `type_args` and
+    /// [`crate::ops::Const`] values for inputs at [`crate::IncomingPort`]s,
+    /// try to evaluate the operation.
+    fn fold(
+        &self,
+        type_args: &[TypeArg],
+        consts: &[(crate::IncomingPort, crate::ops::Const)],
+    ) -> ConstFoldResult;
+}
+
+impl Debug for Box<dyn ConstFold> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "<custom constant folding>")
+    }
+}
+
+impl Default for Box<dyn ConstFold> {
+    fn default() -> Self {
+        Box::new(|&_: &_| None)
+    }
+}
+
+/// Blanket implementation for functions that only require the constants to
+/// evaluate - type arguments are not relevant.
+impl<T> ConstFold for T
+where
+    T: Fn(&[(crate::IncomingPort, crate::ops::Const)]) -> ConstFoldResult + Send + Sync,
+{
+    fn fold(
+        &self,
+        _type_args: &[TypeArg],
+        consts: &[(crate::IncomingPort, crate::ops::Const)],
+    ) -> ConstFoldResult {
+        self(consts)
+    }
+}