feat: Rewrite for inlining a single Call #1934

acl-cqc · 2025-02-25T10:12:31Z

Including adding a new method for copying descendants of another node (preserving edges incoming to that subtree).

The choice of copying descendants rather than a subtree and its root may seem a little strange. The minor difference is that in most usecases (others are peeling tail loops, and monomorphization which can be refactored to use this too) we're gonna update the root node optype too, so we might as well just make the new optype first rather than copy+overwrite. The bigger differences are

Edges between the copied subtree and the root (e.g. recursive Calls) want to point to the original FuncDefn, not the inlined copy (which will be a DFG). This helps the inline-call usecase, and may help monomorphization (recursive calls may have different type args, so we should monomorphize their targets too), and makes no difference for loop peeling (there are no edges to/from a TailLoop from inside).
Edges already existing to the new root (the old Call node) can be left untouched, rather than having to be moved (and the old Call node deleted). Again I think this makes no difference for loop peeling, nor monomorphization.

closes #1833, #1886

codecov · 2025-02-25T10:14:58Z

Codecov Report

Attention: Patch coverage is 85.36585% with 30 lines in your changes missing coverage. Please review.

Project coverage is 83.72%. Comparing base (929edb6) to head (c995322).

Files with missing lines	Patch %	Lines
hugr-core/src/hugr/rewrite/inline_call.rs	86.50%	3 Missing and 19 partials ⚠️
hugr-core/src/hugr/hugrmut.rs	80.95%	7 Missing and 1 partial ⚠️

Additional details and impacted files

@@           Coverage Diff            @@
##             main    #1934    +/-   ##
========================================
  Coverage   83.71%   83.72%            
========================================
  Files         196      197     +1     
  Lines       37533    37738   +205     
  Branches    34346    34551   +205     
========================================
+ Hits        31422    31597   +175     
- Misses       4328     4338    +10     
- Partials     1783     1803    +20

Flag	Coverage Δ
python	`92.34% <ø> (ø)`
rust	`82.93% <85.36%> (+0.01%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

acl-cqc · 2025-02-25T10:23:28Z

hugr-core/src/hugr/hugrmut.rs

@@ -448,6 +456,50 @@ impl<T: RootTagged<RootHandle = Node> + AsMut<Hugr>> HugrMut for T {
        }
        translate_indices(node_map)
    }
+
+    fn copy_descendants(&mut self, root: Node, new_parent: Node) -> HashMap<Node, Node> {
+        // TODO should we check that we will not invalidate the Hugr?


My general feeling is that this is too hard, at the least, it involves at least O(size of subtree) work (probably * O(depth)), never mind Dom edges; and it'd only be meaningful if the new location was already in its final resting place in the hierarchy anyway. But it does leave a bit of a bad taste that a "valid" Rewrite could make a valid Hugr invalidate without warning or any method of detection....

Any linear-typed edge incoming to root

I don't think this a problem on its face. the ports of root have nothing to do with this. The copied Input node will have a linear outgoing edge, but it will be connected exactly once. A bigger problem is that new_parent will now have two input nodes, unless you carefully arrange that it is initially empty

Non-local edges can't happen for FuncDefns, so that's nice.

Other methods on HugrMut can invalidate the hugr, e.g. disconnect. @aborgna-q What do you think? Should this method be on HugrMutInternals?.

I would prefer this comment(well, at least the non-local edges bit) to be a doc on the definition in the trait.

Yeah, I think the linear-typed edge comment only applied to when I was copying subtrees not descendants, so I deleted that and left the rest....

Your example of disconnect is good, so if the best we can do is to make this a doc comment (and raise an issue for better handling/detection of "nonlocal edges incoming to a subtree", if we don't have one already...hmmm), then let's do that....right?

I would say yes. I tagged Agustin because he asked me to move a "dangerous" function, and he may disagree.

@doug-q's comment seems reasonable imho.

We should describe how this may break the structure in the description, including the doubled i/o nodes.

acl-cqc · 2025-02-25T10:55:43Z

Coverage looks poor because of use of ? in tests rather than .unwrap(), the ? is considered only partial because these things never error. I can make the tests more verbose but it won't make them easier to read...

ss2165 · 2025-02-25T11:42:06Z

hugr-core/src/hugr/hugrmut.rs

@@ -146,6 +146,14 @@ pub trait HugrMut: HugrMutInternals {
        self.hugr_mut().remove_node(node);
    }

+    /// Copies the strict descendants of `root` to under the `new_parent`.
+    /// (That is, the immediate children of root, are copied to make children of `new_parent`).
+    fn copy_descendants(&mut self, root: Node, new_parent: Node) -> HashMap<Node, Node> {


could a test be refactored to use this function so it is covered?

hugr-core/src/hugr/hugrmut.rs

ss2165 · 2025-02-25T11:52:03Z

hugr-core/src/hugr/hugrmut.rs

+        // Copy the optypes, metadata, and hierarchy
+        for (&node, &new_node) in node_map.iter() {
+            for ch in self.children(node).collect::<Vec<_>>() {
+                self.set_parent(*node_map.get(&ch).unwrap(), new_node);


looks like the hierarchy part isn't covered, do you need a test with more nesting?

ss2165 · 2025-02-25T11:59:14Z

hugr-core/src/hugr/rewrite/inline_call.rs

+            let [i] = fb.input_wires_arr();
+            let add = fb.add_dataflow_op(IntOpDef::iadd.with_log_width(5), [i, cst1])?;
+            let call = fb.call(
+                &FuncID::<true>::from(fb.container_node()),


not for this PR but this recursive handle should probably in the function builder API

ss2165 · 2025-02-25T12:01:08Z

hugr-core/src/hugr/rewrite/inline_call.rs

+            }
+            assert_eq!(ancestors.next(), Some(main.node()));
+            assert_eq!(ancestors.next(), Some(hugr.root()));
+            assert_eq!(ancestors.next(), None);


doug-q · 2025-02-26T08:24:33Z

hugr-core/src/hugr/hugrmut.rs

@@ -146,6 +146,14 @@ pub trait HugrMut: HugrMutInternals {
        self.hugr_mut().remove_node(node);
    }

+    /// Copies the strict descendants of `root` to under the `new_parent`.
+    /// (That is, the immediate children of root, are copied to make children of `new_parent`).
+    fn copy_descendants(&mut self, root: Node, new_parent: Node) -> HashMap<Node, Node> {


Consider returning a BTreeMap over a HashMap here, as a small help for determinism.

doug-q · 2025-02-26T08:34:40Z

hugr-core/src/hugr/hugrmut.rs

+        }*/
+        let mut nodes = Vec::new();
+        let mut q = VecDeque::from_iter(self.children(root));
+        while let Some(n) = q.pop_front() {


Can you use DescendentsView here? This is such a common operation I think we should provide a helper for it. I have code golfed

iter::successors(Some(vec![root]), |nodes| { let children = nodes.iter().flat_map(|&n| hugr.children(n)).collect_vec(); (!children.is_empty()).then_some(children) })

I originally avoided this because I didn't want to "build a whole DescendantsGraph" but now I realize that a DescendantsGraph isn't much more than the collection of nodes that we want.

However, in fact there is a method in Hierarchy that does pretty much exactly what we want. So I've used that. Do you (still) think it's worth adding a helper to HugrView??

In general yes, for you here, no. Hierarchy is not available out-of-crate.

doug-q · 2025-02-26T08:38:38Z

hugr-core/src/hugr/hugrmut.rs

+        // Copy the optypes, metadata, and hierarchy
+        for (&node, &new_node) in node_map.iter() {
+            for ch in self.children(node).collect::<Vec<_>>() {
+                self.set_parent(*node_map.get(&ch).unwrap(), new_node);


This is nondeterministic because you are iterating a HashMap. the order of children in their parent will depend on iteration order. It's a little out of scope for this PR, but would you consider changing translate_indices to return a BTreeMap?

EDIT: no I'm wrong. You mutate the nodes in a nondeterministic order, but I don't think this is actually nondeterministic. for each parent you set_parent it's children in a deterministic order.

Yes, ISTR this hit me just getting Input/Output in the wrong order, so I had to make sure it did it "right"...

doug-q · 2025-02-26T09:01:09Z

hugr-core/src/hugr/rewrite/inline_call.rs

@@ -0,0 +1,234 @@
+//! Rewrite to inline a Call to a FuncDefn by copying the body of the function
+//! into a DFG which replaces the Call node.
+use thiserror::Error;


Can you use derive_more::Error? One day we will be able to delete the thiserror dependency

hugr-core/src/hugr/rewrite/inline_call.rs

doug-q · 2025-02-26T09:04:20Z

hugr-core/src/hugr/rewrite/inline_call.rs

+        Ok(())
+    }
+
+    fn apply(self, h: &mut impl HugrMut) -> Result<(), Self::Error> {


I don't think you are considering the case of polymorphic functions here? I think you need to substitute the type args of the call into the signature and all descendents.

I fo think it's ok to reject polymorphic functions for now

Oh dang that's a good point. Knew some things would get forgotten what with this sleeping over hack fortnight...

So I could do this later, but it's not much new code - you just need to iterate through the same nodes and apply the substitution. However, this does involve cloning the OpType in copy_descendants followed by applying the substitution (which is kinda a copy itself). So we could

decide this bit of extra copying isn't important

make substitution destructive (mutating). This would be perhaps the biggest change.

pass a impl Fn(&OpType) -> OpType into copy_descendants. Seems a little bit arbitrary (shouldn't we also have callbacks to copy metadata, etc....) but I think it's reasonable to argue that we should have this one. What do you think?

In fact on the third I am tempted to make copy_descendants take an Option<Substitution> rather than a generic callback. Substitution is pretty fundamental and it avoids the argument that we should also do the same thing for e.g. metadata ;)

I don't think the copying is important.
I do like passing an optional substitution, this is very nice.
It seems like a bit of a weird inter-module dependency, this is very low level, substitution pretty high level. But this is just a vibe.

acl-cqc added 12 commits February 25, 2025 10:05

Add InlineCall - v0 supporting recursion but no other incoming edges

7a3ac0b

fix verify to also check call target

6d8dfc1

Update portgraph to 0.13.2

0d5b5af

Add HugrMut::copy_subtree. TODO check for invalidation?

75f8842

InlineCall uses copy_subtree

4c345eb

First test; fix copying of optypes,metadata,hierarchy

120f17c

WIP test

c10c81a

Fix recursion by changing copy_subtree => copy_descendants

45dc680

Validate in test

d3f9c62

Tidy tests w/helper methods calls/extension_ops

231b125

refactor test: inline get_nonrec_call

756902a

Fix tests for --all-features

8db58de

Update that comment

45dd9ac

acl-cqc commented Feb 25, 2025

View reviewed changes

test errors, improve+derive InlineCallError

c995322

acl-cqc marked this pull request as ready for review February 25, 2025 10:54

acl-cqc requested a review from a team as a code owner February 25, 2025 10:54

acl-cqc requested a review from ss2165 February 25, 2025 10:54

ss2165 reviewed Feb 25, 2025

View reviewed changes

doug-q reviewed Feb 26, 2025

View reviewed changes

doug-q mentioned this pull request Feb 25, 2025

Inline hugr functions CQCL/hugr-qir#12

Open

acl-cqc and others added 3 commits February 27, 2025 09:45

Use hierarchy.descendants

5a1f2f8

oops, fix prev

7f86624

Add InlineCall::new

b32065b

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Rewrite for inlining a single Call #1934

feat: Rewrite for inlining a single Call #1934

acl-cqc commented Feb 25, 2025 •

edited

Loading

codecov bot commented Feb 25, 2025 •

edited

Loading

acl-cqc Feb 25, 2025

doug-q Feb 26, 2025

acl-cqc Feb 27, 2025

doug-q Feb 27, 2025

aborgna-q Feb 27, 2025

acl-cqc commented Feb 25, 2025

ss2165 Feb 25, 2025

ss2165 Feb 25, 2025

ss2165 Feb 25, 2025

ss2165 Feb 25, 2025

doug-q Feb 26, 2025

doug-q Feb 26, 2025

acl-cqc Feb 27, 2025

doug-q Feb 27, 2025

doug-q Feb 26, 2025 •

edited

Loading

acl-cqc Feb 27, 2025

doug-q Feb 26, 2025

doug-q Feb 26, 2025

acl-cqc Feb 27, 2025

acl-cqc Feb 27, 2025

doug-q Feb 27, 2025

feat: Rewrite for inlining a single Call #1934

Are you sure you want to change the base?

feat: Rewrite for inlining a single Call #1934

Conversation

acl-cqc commented Feb 25, 2025 • edited Loading

codecov bot commented Feb 25, 2025 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

acl-cqc commented Feb 25, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

doug-q Feb 26, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

acl-cqc commented Feb 25, 2025 •

edited

Loading

codecov bot commented Feb 25, 2025 •

edited

Loading

doug-q Feb 26, 2025 •

edited

Loading