hydro_lang/rewrites/
properties.rs

Help
1use std::collections::HashSet;
2
3use stageleft::*;
4
5use crate::ir::{HydroLeaf, HydroNode, transform_bottom_up};
6
7/// Structure for tracking expressions known to have particular algebraic properties.
8///
9/// # Schema
10///
11/// Each field in this struct corresponds to an algebraic property, and contains the list of
12/// expressions that satisfy the property. Currently only `commutative`.
13///
14/// # Interface
15///
16/// "Tag" an expression with a property and it will add it to that table. For example, [`Self::add_commutative_tag`].
17/// Can also run a check to see if an expression satisfies a property.
18#[derive(Default)]
19pub struct PropertyDatabase {
20    commutative: HashSet<syn::Expr>,
21}
22
23impl PropertyDatabase {
24    /// Tags the expression as commutative.
25    pub fn add_commutative_tag<
26        'a,
27        I,
28        A,
29        F: Fn(&mut A, I),
30        Ctx,
31        Q: QuotedWithContext<'a, F, Ctx> + Clone,
32    >(
33        &mut self,
34        expr: Q,
35        ctx: &Ctx,
36    ) -> Q {
37        let expr_clone = expr.clone();
38        self.commutative.insert(expr_clone.splice_untyped_ctx(ctx));
39        expr
40    }
41
42    pub fn is_tagged_commutative(&self, expr: &syn::Expr) -> bool {
43        self.commutative.contains(expr)
44    }
45}
46
47// Dataflow graph optimization rewrite rules based on algebraic property tags
48// TODO add a test that verifies the space of possible graphs after rewrites is correct for each property
49
50fn properties_optimize_node(node: &mut HydroNode, db: &mut PropertyDatabase) {
51    match node {
52        HydroNode::ReduceKeyed { f, .. } if db.is_tagged_commutative(&f.0) => {
53            dbg!("IDENTIFIED COMMUTATIVE OPTIMIZATION for {:?}", &f);
54        }
55        _ => {}
56    }
57}
58
59pub fn properties_optimize(ir: &mut [HydroLeaf], db: &mut PropertyDatabase) {
60    transform_bottom_up(ir, &mut |_| (), &mut |node| {
61        properties_optimize_node(node, db)
62    });
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68    use crate::FlowBuilder;
69    use crate::deploy::SingleProcessGraph;
70    use crate::location::Location;
71
72    #[test]
73    fn test_property_database() {
74        let mut db = PropertyDatabase::default();
75
76        assert!(
77            !db.is_tagged_commutative(&(q!(|a: &mut i32, b: i32| *a += b).splice_untyped_ctx(&())))
78        );
79
80        let _ = db.add_commutative_tag(q!(|a: &mut i32, b: i32| *a += b), &());
81
82        assert!(
83            db.is_tagged_commutative(&(q!(|a: &mut i32, b: i32| *a += b).splice_untyped_ctx(&())))
84        );
85    }
86
87    #[test]
88    fn test_property_optimized() {
89        let flow = FlowBuilder::new();
90        let mut database = PropertyDatabase::default();
91
92        let process = flow.process::<()>();
93        let tick = process.tick();
94
95        let counter_func = q!(|count: &mut i32, _| *count += 1);
96        let _ = database.add_commutative_tag(counter_func, &tick);
97
98        unsafe {
99            process
100                .source_iter(q!(vec![]))
101                .map(q!(|string: String| (string, ())))
102                .tick_batch(&tick)
103        }
104        .fold_keyed(q!(|| 0), counter_func)
105        .all_ticks()
106        .for_each(q!(|(string, count)| println!("{}: {}", string, count)));
107
108        let built = flow
109            .optimize_with(|ir| properties_optimize(ir, &mut database))
110            .with_default_optimize::<SingleProcessGraph>();
111
112        insta::assert_debug_snapshot!(built.ir());
113
114        let _ = built.compile_no_network();
115    }
116}
hydro_lang/rewrites/properties.rs

hydro_lang/rewrites/
properties.rs
