Flatten regions #1347

lukemaurer · 2023-05-03T12:07:46Z

It's very important to take

let x =
  let y = E in
  F
in
G

and rewrite it to

let y = E in
let x = F in
G

so that, for instance, when simplifying G we're free to make use of y. This can be crucial for eliminating allocations:

let x =
  let y = E in
  (y, y)
in
fst x

Here we can't easily eliminate the pair because we don't have a binding for the result of the projection. If we lift first, however:

let y = E in
let x = (y, y) in
fst x

This trivially reduces to let y = E in y (and further to E, as it happens).

Unfortunately, regions introduced by inlined calls get in the way of lifting:

let x =
  region (
    let y = E in
    F
  )
in
G

In this case, currently we give up and don't allow y to be visible to G at all. It is tempting to write:

region (
  let y = E in
  let x = F in
  G
)

Unfortunately, this moves G into the region, which is an invalid transformation. However, we can fix this with an exclave:

region (
  let y = E in
  let x = F in
  exclave (
    G
  )
)

This patch enhances the let-floating transform in flambda1 by lifting most regions out in this way. This works whenever there isn't already an exclave (unless the exclave is particularly convenient - in particular, not buried in an if or for or what-have-you) and will hopefully eliminate some heap allocations that have cropped up when turning on local allocations.

This squashes 'flatten-regions' into a single commit over 'remove-regions-with-exclaves' and 'cmm-nested-exclaves' (PRs ocaml-flambda#1524 and ocaml-flambda#1529, respectively). Squashed commit of the following: commit bd23d8bd75478abafb67f0e791ef77bbd8764b22 Author: Luke Maurer <lmaurer@janestreet.com> Date: Fri Jun 30 16:03:06 2023 +0100 Clean up code commit 0ac81ecfc5b84eb286cbd5f0afc66d7f73b3eb92 Merge: c3a9203f9 9449e65 Author: Luke Maurer <lmaurer@janestreet.com> Date: Fri Jun 30 14:59:27 2023 +0100 Merge branch 'cmm-nested-exclaves' into flatten-regions commit c3a9203f9de900d8acbe77b781debbfb686a3d52 Merge: c1b79d7 d617694 Author: Luke Maurer <lmaurer@janestreet.com> Date: Fri Jun 30 14:57:23 2023 +0100 Merge branch 'remove-regions-with-exclaves' into flatten-regions commit c1b79d7 Author: Luke Maurer <lmaurer@janestreet.com> Date: Fri Jun 30 14:49:07 2023 +0100 Add example commit 033294f Author: Luke Maurer <lmaurer@janestreet.com> Date: Thu Jun 8 11:21:22 2023 +0100 Code review and provisional fixes These are fixes to issues that are really second- or third-order effects of region lifting and probably need their own PRs, but I'm committing them here for testing purposes. commit 65ad3aa Merge: e88df82 d65f752 Author: Luke Maurer <lmaurer@janestreet.com> Date: Thu May 25 17:00:33 2023 +0100 Merge remote-tracking branch 'upstream/main' into flatten-regions commit e88df82 Merge: de8cf43 9d3c1c1 Author: Luke Maurer <lmaurer@janestreet.com> Date: Thu May 11 13:21:45 2023 +0100 Merge remote-tracking branch 'upstream/main' into flatten-regions commit de8cf43 Author: Luke Maurer <lmaurer@janestreet.com> Date: Wed May 10 16:47:44 2023 +0100 Follow renaming of `Tail` to `Exclave` after merge commit ceae59a Merge: 8a717aa 63997c9 Author: Luke Maurer <lmaurer@janestreet.com> Date: Wed May 10 16:37:23 2023 +0100 Merge remote-tracking branch 'upstream/main' into flatten-regions commit 8a717aa Author: Luke Maurer <lmaurer@janestreet.com> Date: Wed May 10 16:26:33 2023 +0100 Don't try to be clever with `let x = M in tail x` There's nothing actually wrong with `let x = M in tail x`, so don't try to reduce it to just `M`. This works in that narrow case, but the code that was doing this transformation didn't notice if the body of the tail is more than just `x`. Since the transformation doesn't actually gain anything, better to be rid of it than to make things more complicated trying to get it right. Also updated a few comments and added an invariant check. commit 6fa1ce6 Author: Luke Maurer <lmaurer@janestreet.com> Date: Fri Apr 21 15:58:07 2023 +0100 Make `map_region_tail` look through regions as well commit d9043c9 Author: Luke Maurer <lmaurer@janestreet.com> Date: Fri Apr 21 14:14:19 2023 +0100 Check for close-on-apply on method calls as well as functions commit 58bf88a Author: Luke Maurer <lmaurer@janestreet.com> Date: Wed Apr 19 17:05:31 2023 +0100 Fix `Cmmgen` `Cmmgen` was (hopefully) the last remaining place where things were broken by the assumption that `region (region (tail (tail ...)))` never happens. commit 66132fb Author: Luke Maurer <lmaurer@janestreet.com> Date: Tue Apr 18 18:43:08 2023 +0100 Don't try to handle regions in `Flambda.fold_lets_option` It used to be much more necessary, but no more, and it interferes with the handling of regions using `enter_region`. commit 907627d Author: Luke Maurer <lmaurer@janestreet.com> Date: Tue Apr 18 18:41:18 2023 +0100 Refactor and clean up a bit commit dccb971 Author: Luke Maurer <lmaurer@janestreet.com> Date: Tue Apr 18 17:55:02 2023 +0100 Rewrite `let x = M in tail x` as `M` This should get rid of some annoying instances of extra variables introduced just to lift, and in particular stop making tail calls into non-tail calls. commit 6126e9e Merge: 5b62db1 e138734 Author: Luke Maurer <lmaurer@janestreet.com> Date: Wed Apr 12 14:31:25 2023 +0100 Merge remote-tracking branch 'upstream/main' into flatten-regions commit 5b62db1 Author: Luke Maurer <lmaurer@janestreet.com> Date: Wed Nov 30 17:37:40 2022 +0000 Lift regions along with `let`s This transforms expressions of the form ``` let x = region ( let y = E1 in E2 ) in E3 ``` into ``` region ( let y = E1 in let x = E2 in tail ( E3 ) ) ``` so that `y` is available when simplifying `E3` (as is the advantage of lifting `let`s to begin with). This requires that `E3` not already have a `tail` expression (or tail call) inside an `if` branch or any other construct besides the body of a `let`, `let mutable`, `region`, or `tail`.

middle_end/flambda/lift_code.ml

stedolan · 2023-08-15T12:47:26Z

middle_end/flambda/lift_code.mli

+(** Lift lets in an expression. Note that [toplevel] here means either truly
+    toplevel _or_ inside a toplevel function body. *)
+val lift_lets_expr
+   : Flambda.t -> toplevel:bool -> in_closure:yes_no_or_maybe -> Flambda.t


I think it might be cleaner if this was may_lift_regions:bool instead of in_closure:yes_no_or_maybe, since the No and Maybe cases aren't distinguished by this code and it's not obvious why being in a closure matters here.

I find it easier to keep track of simple syntactic properties like “Is it in a closure?” than “Are we allowed to lift a region here?”, especially when it comes to lines like this:

flambda-backend/middle_end/flambda/lift_code.ml

Line 334 in 9b05a30

~f:(lift_lets_expr ~toplevel ~in_closure:Yes) set))

The answer to “Is it in a closure?” is clearly just “yes,” whereas the answer to “Are we allowed to lift a region here?” is contingent on flambda1's quirks. This way we get to express our policy about when to lift regions exactly once, in should_lift_regions.

It was at one point a boolean, but I got tired of getting mixed up between “not known to be true” and “known not to be true” and decided to isolate that concern as well.

stedolan · 2023-08-15T12:48:56Z

middle_end/flambda/lift_code.ml

-  extract acc (W.of_expr body)
+  extract acc dest (W.of_expr body) ~in_closure
+
+(* C[let x = region M in x] *)


I don't think the body of the let is constrained to be x here?

Ah, this isn't very clear: body is the body of the region, not the body of the let. So it's M in the expression.

Would these comments be any clearer if they read
<acc>[let <dest> = region <body> in <dest>] and such? I find it helpful to be able to map a notional equational theory onto the implementation code, but clearly I went too far and tried to write a POPL paper in the comments ...

stedolan · 2023-08-15T12:50:29Z

middle_end/flambda/lift_code.ml

+       will move arbitrary computations - if there is a [Region] but no
+       [Exclave], this means we're moving those computations into a different
+       region. It may be that [acc_expr] already has an [Exclave] (because we
+       lifted it out of [body]), but otherwise we need to add it. *)


Is a single Exclave definitely enough to fix this? Could there be two Regions?

There's only one Region we need to worry about, and that's the one we're adding (on line 242). The extreme thing that acc_expr could do is close more regions (that is, have more Exclaves), which will be handled by outer calls to extract_region.

mshinwell · 2023-12-04T17:36:08Z

We don't presently need this now we are all on flambda2, but if locals ends up upstream before flambda2, it would probably be needed.

mshinwell added the middle end label May 4, 2023

lukemaurer marked this pull request as ready for review May 11, 2023 12:22

lukemaurer requested review from mshinwell, gretay-js and xclerc as code owners May 11, 2023 12:22

lukemaurer force-pushed the flatten-regions branch from f9bcebc to c1b79d7 Compare June 30, 2023 13:52

lukemaurer requested a review from lthls as a code owner June 30, 2023 13:52

lukemaurer force-pushed the flatten-regions branch from c1b79d7 to 7691e17 Compare June 30, 2023 15:12

lukemaurer force-pushed the flatten-regions branch from 7691e17 to 9b05a30 Compare August 9, 2023 14:25

stedolan reviewed Aug 15, 2023

View reviewed changes

mshinwell closed this Dec 4, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Flatten regions #1347

Flatten regions #1347

lukemaurer commented May 3, 2023 •

edited

Loading

stedolan Aug 15, 2023

lukemaurer Aug 22, 2023

stedolan Aug 15, 2023

lukemaurer Aug 22, 2023

stedolan Aug 15, 2023

lukemaurer Aug 22, 2023

mshinwell commented Dec 4, 2023

Flatten regions #1347

Flatten regions #1347

Conversation

lukemaurer commented May 3, 2023 • edited Loading

stedolan Aug 15, 2023

Choose a reason for hiding this comment

lukemaurer Aug 22, 2023

Choose a reason for hiding this comment

stedolan Aug 15, 2023

Choose a reason for hiding this comment

lukemaurer Aug 22, 2023

Choose a reason for hiding this comment

stedolan Aug 15, 2023

Choose a reason for hiding this comment

lukemaurer Aug 22, 2023

Choose a reason for hiding this comment

mshinwell commented Dec 4, 2023

lukemaurer commented May 3, 2023 •

edited

Loading