tket2 icon indicating copy to clipboard operation
tket2 copied to clipboard

BorrowSquash pass to reorder and remove redundant borrows on borrow_arrays

Open lmondada opened this issue 4 months ago • 7 comments

Here's an overview of a Hugr->Hugr transform before we encode subcircuits for pytket:

  1. Read-only analysis pass to find the borrow interval (i.e. pair up borrows with returns)
    • see https://github.com/CQCL/tket2/pull/1072
    • hoping to pass over to @tatiana-s to test many cases
  2. Reorder intervals. Establish constraints between borrow intervals and use toposort to find order that i) satisfies all constraints and ii) groups borrows of the same resource together
    • @lmondada will work on this next
    • I have specified the algorithm in #1074
  3. Remove redundant pairs of return-borrows that borrow the same resource with the same borrow index
    • the redundant pairs will be guaranteed to be next to eachother, so this should be as simple as creating rewrites that merge two ops

lmondada avatar Aug 26 '25 07:08 lmondada

We might be able to close #1074 (and include it here) as (1) I think I have a simpler way to do the reordering, (2) the interface between analysis/transform is more complex than "return a list of all intervals".

  • The latter because, so long as all intervals are consistent for lifetime of one Resource/array, there is no reason not to squash return/borrows for that array even if other arrays fail
  • We need to be careful about lifetime of an array, i.e. resource tracking. IIUC the semantics of borrow/release mean (@mark-koch please confirm) that there should be no elements borrowed during any other array op (e.g. pop, swap with external element, passing the whole array to a function - the borrowed elements fine). In which case we can say any such op consumes the array Resource and returns a new Resource. (Then we expect borrow/return to match up pairwise for any resource.)
    • A tricky case is nested arrays - here borrowing from the parent array creates a new Resource for the inner array, which may itself have elements borrowed/returned; at some point the inner array or some other is returned to the outer array.
    • This is not what's currently in the PR #1072 !
    • Also we need analysis to fail at the end of the resource is there's anything still borrowed

Following all that, we now have that analysis for an array produces a list of borrow+return intervals (that match up), and there are no other ops. Assuming for now that all indices are static, elision/squashing can then proceed (for that array) by processing the borrows/returns in sequence as follows:

  • Maintain the last_outport producing "the current state of the array" (bitmask) - at first this is the Array::new or Input node etc.
  • Also borrowed a map from array indices (that have been borrowed) to the borrow node, and Option. Initially empty.
  • Consider each borrow/return node in order; the analysis ensures only three cases can occur:
    • For a borrow node, where the index not has been seen - record in borrowed, and emit it (set its array input to the last_outport then set last_outport to its array output).
    • For a return node, (analysis ensures) the index has been seen, and there will be a borrow node (and no previous return). Just record the new return node in borrowed.
    • For a borrow node where the index has been seen, analysis ensures there will be a return node recorded in borrowed. Remove the old borrow and the new return; wire the return's value input to the borrow's value output. Remove the return node from borrowed (leaving just the earlier/original borrow node).
  • Once we've done that, borrowed contains only pairs of borrow+return. Emit the returns (by wiring last_outport to their array input, and setting last_outport to their array output).

acl-cqc avatar Oct 02 '25 18:10 acl-cqc

I agree this looks much simpler!

Just record the new return node in borrowed.

Where exactly do you record this? Is that the Option that you mention within a borrowed entry?

The fact that no operations other than borrow/returns will be found on the array wires indeed simplifies the logic greatly. Have you decided what to do with dynamic indices (and exception semantics)? It might be in that case that you need to "flush" all borrows at that point, but that should not be too hard.

lmondada avatar Oct 06 '25 03:10 lmondada

Just record the new return node in borrowed.

Where exactly do you record this? Is that the Option that you mention within a borrowed entry?

Yes, so there's a map from element index to borrow node and Option of return; the latter gets filled in. If there's another borrow, the Option is emptied (and return->borrow is elided). At the end, the Options are all full (all borrows are returned - the analysis says, don't run if this isn't the case!), so again we only have borrow-return pairs (even if not the original pairs), and we output those.

Have you decided what to do with dynamic indices (and exception semantics)? It might be in that case that you need to "flush" all borrows at that point, but that should not be too hard.

Exception semantics don't really matter if we are only eliding return-borrow, 'coz the index for the (second) borrow has already been checked. Yes, except where there are dynamic indices...

So I am not planning on allowing dynamic indices any time soon TBH, but I've found two stages

  • analysis stage checks borrows/returns are correctly paired
  • transform stage looks for return-borrow pairs, also in an extension looks for borrow followed by corresponding return of the same value - in which case we can elide that.

TBH I've found the transform stage quite hard to break up, I mean there could be a separate "find the pairs to elide" step that doesn't do any mutation, but when we come to borrow-return of the same value, we want to have been doing the mutation as we go along so it's easy to look for the value being the same. So I've kept everything there internal to one method, whose signature stays the same, so updates like dynamic indexing etc. can follow later (as non-breaking releases).

However, something like you say, yes....no borrow-return elision if there's any dynamically-indexed borrows made between them; but you have to keep modelling the state of what (else) is borrowed even though you're not doing the elision. I guess there'll be a new step of seeing a return and clearing it (and the preceding borrow) out of the map.

acl-cqc avatar Oct 06 '25 13:10 acl-cqc

IIUC the semantics of borrow/release mean (@mark-koch please confirm) that there should be no elements borrowed during any other array op (e.g. pop, swap with external element, passing the whole array to a function - the borrowed elements fine).

This is currently (mostly) the case for Hugrs coming from Guppy, however I think does not generally hold for all Hugrs: The runtime checks guarantee that that the elements affected by an op (e.g. the popped element) are not borrowed, however we don't check the other elements.

There is also the following situation in Guppy when using nested arrays:

cx(qs[0][0], qs[0][1])

# Turns into:
a = borrow(qs, 0)
b = borrow(a, 0)
return(qs, 0, a)
c = borrow(qs, 0)
d = borrow(c, 1)
return(qs, 0, c)
cx(b, d)
e = borrow(qs, 0)
return(e, 0, b)
return(qs, 0, e)
f = borrow(qs, 0)
return(f, 1, d)
return(qs, 0, f)

Here, we return arrays with holes into the outer array.

Finally, note that there are plans to add functions to Guppy to expose the notion of "taking elements out of arrays and leaving a hold" (see https://github.com/CQCL/guppylang/pull/1165). If we add that, then we can no longer make any assumption on which elements are borrowed vs not

mark-koch avatar Oct 06 '25 14:10 mark-koch

Thanks @mark-koch !

I think we can get a fair way in the analysis by having a notion of an op that "clobbers" a specific index of the bitmask (requiring the correct state of that bitmask element to be restored before the op, and starting afresh with no knowledge of what's in that element afterwards). I mean, ok we can get further still with more detailed modelling of more ops, but that'd be a good start.

However, it also sounds like we want the elision to work in cases where the borrows/returns are not all correctly paired, which kinda means the first analysis step can't really do much of anything.

Your "returning arrays with holes into the outer array" is more problematic still. Current use of ResourceScope means we don't realize that borrowing an index, then returning something, then borrowing it again, is actually all "the same value". This will be hopefully be revealed by return-borrow elision, but that happens later, after we've finished with the ResourceScope. It's hard to see how this isn't a case of "rerun the transformation if it did anything", which is not great :(....so may need to abandon ResourceScope.

I was thinking perhaps we could change the guppy compiler so that we don't return the outer array until after returning the inner arary, but this isn't correct if both arguments to CX borrow the same element of the outer array (perhaps with different variables!)...argh.

acl-cqc avatar Oct 06 '25 20:10 acl-cqc

I think we can get there if we don't use ResourceScope but do the analysis sufficiently incrementally (one wire at a time, in a carefully controlled order!). Much of the same code, but more intricately interwoven.

The ops in https://github.com/CQCL/guppylang/pull/1165 are basically just user-callable borrow and return, no? (So, not paired up!)

acl-cqc avatar Oct 07 '25 07:10 acl-cqc

Mark's nested array example is now addressed in the PR, and the future get/put ops can be addressed (in future) via #1176

acl-cqc avatar Oct 14 '25 18:10 acl-cqc