Merge branch 'profiler' of https://github.com/davidchuyaya/hydroflow …

…into profiler
davidchuyaya · Mar 14, 2024 · 004a718 · 004a718
2 parents 7ac5fe7 + 8043c36
commit 004a718
Show file tree

Hide file tree

Showing 5 changed files with 59 additions and 4 deletions.
diff --git a/hydroflow_plus/src/ir.rs b/hydroflow_plus/src/ir.rs
@@ -79,7 +79,7 @@ pub enum HfPlusLeaf {
 impl HfPlusLeaf {
     pub fn transform_children(
         self,
-        transform: impl Fn(HfPlusNode, &mut SeenTees) -> HfPlusNode,
+        mut transform: impl FnMut(HfPlusNode, &mut SeenTees) -> HfPlusNode,
         seen_tees: &mut SeenTees,
     ) -> HfPlusLeaf {
         match self {
@@ -215,7 +215,7 @@ pub type SeenTees = HashMap<*const RefCell<HfPlusNode>, Rc<RefCell<HfPlusNode>>>
 impl HfPlusNode {
     pub fn transform_children(
         self,
-        transform: impl Fn(HfPlusNode, &mut SeenTees) -> HfPlusNode,
+        mut transform: impl FnMut(HfPlusNode, &mut SeenTees) -> HfPlusNode,
         seen_tees: &mut SeenTees,
     ) -> HfPlusNode {
         match self {

diff --git a/hydroflow_plus/src/lib.rs b/hydroflow_plus/src/lib.rs
@@ -34,6 +34,7 @@ pub use builder::FlowBuilder;
 pub mod ir;
 
 pub mod persist_pushdown;
+pub mod profiler;
 
 #[derive(Clone)]
 pub struct RuntimeContext<'a> {

diff --git a/hydroflow_plus/src/profiler.rs b/hydroflow_plus/src/profiler.rs
@@ -38,7 +38,7 @@ pub fn profiling<'a>(ir: Vec<HfPlusLeaf>, context: RuntimeContext<'a>, counters:
 #[cfg(test)]
 mod tests {
     use stageleft::*;
-
+    use std::cell::RefCell;
     use crate::{Location, MultiGraph};
 
     #[test]
@@ -56,7 +56,8 @@ mod tests {
 
         insta::assert_debug_snapshot!(&built.ir);
 
-        let pushed_down = built.optimize_with(|ir| super::profiling(ir, flow.runtime_context()));
+        let counters = RuntimeData::new("Fake");
+        let pushed_down = built.optimize_with(|ir| super::profiling(ir, flow.runtime_context(), counters));
 
         insta::assert_debug_snapshot!(&pushed_down.ir);
     }

diff --git a/.../src/snapshots/hydroflow_plus__profiler__tests__predicate_pushdown_through_map-2.snap.new b/.../src/snapshots/hydroflow_plus__profiler__tests__predicate_pushdown_through_map-2.snap.new
@@ -0,0 +1,31 @@
+---
+source: hydroflow_plus/src/profiler.rs
+assertion_line: 62
+expression: "&pushed_down.ir"
+---
+[
+    ForEach {
+        f: { use crate :: __staged :: profiler :: tests :: * ; | n | println ! ("{}" , n) },
+        input: Map {
+            f: { use crate :: __staged :: profiler :: * ; let counters_copy = Fake ; let my_id = 0u32 ; let context = & context ; let my_id = 32710u32 ; let counters = Fake ; let my_id = 32710u32 ; { counters_copy . borrow_mut () [my_id as usize] = 0 ; move | v | { hydroflow_plus :: profiler :: increment_counter (context . current_tick () , my_id , & mut counters . borrow_mut () [my_id as usize]) ; v } } },
+            input: Map {
+                f: { use crate :: __staged :: profiler :: tests :: * ; | v | v + 1 },
+                input: Map {
+                    f: { use crate :: __staged :: profiler :: * ; let counters_copy = Fake ; let my_id = 1u32 ; let context = & context ; let my_id = 32710u32 ; let counters = Fake ; let my_id = 32710u32 ; { counters_copy . borrow_mut () [my_id as usize] = 0 ; move | v | { hydroflow_plus :: profiler :: increment_counter (context . current_tick () , my_id , & mut counters . borrow_mut () [my_id as usize]) ; v } } },
+                    input: Persist(
+                        Map {
+                            f: { use crate :: __staged :: profiler :: * ; let counters_copy = Fake ; let my_id = 2u32 ; let context = & context ; let my_id = 32710u32 ; let counters = Fake ; let my_id = 32710u32 ; { counters_copy . borrow_mut () [my_id as usize] = 0 ; move | v | { hydroflow_plus :: profiler :: increment_counter (context . current_tick () , my_id , & mut counters . borrow_mut () [my_id as usize]) ; v } } },
+                            input: Source {
+                                source: Iter(
+                                    { use crate :: __staged :: profiler :: tests :: * ; 0 .. 10 },
+                                ),
+                                produces_delta: false,
+                                location_id: 0,
+                            },
+                        },
+                    ),
+                },
+            },
+        },
+    },
+]
diff --git a/...us/src/snapshots/hydroflow_plus__profiler__tests__predicate_pushdown_through_map.snap.new b/...us/src/snapshots/hydroflow_plus__profiler__tests__predicate_pushdown_through_map.snap.new
@@ -0,0 +1,22 @@
+---
+source: hydroflow_plus/src/profiler.rs
+assertion_line: 57
+expression: "&built.ir"
+---
+[
+    ForEach {
+        f: { use crate :: __staged :: profiler :: tests :: * ; | n | println ! ("{}" , n) },
+        input: Map {
+            f: { use crate :: __staged :: profiler :: tests :: * ; | v | v + 1 },
+            input: Persist(
+                Source {
+                    source: Iter(
+                        { use crate :: __staged :: profiler :: tests :: * ; 0 .. 10 },
+                    ),
+                    produces_delta: false,
+                    location_id: 0,
+                },
+            ),
+        },
+    },
+]