Consider these two somewhat artificial programs (#22404)
Program (P1) Program (P2)
-------------------------------------
let v = <small thunk> in let v = <small thunk> in
join j = case v of (a,b) -> a
in case x of in case x of
A -> case v of (a,b) -> a A -> j
B -> case v of (a,b) -> a B -> j
C -> case v of (a,b) -> b C -> case v of (a,b) -> b
D -> [] D -> []
In (P1), `v` gets allocated, as a thunk, every time this code is executed. But
notice that `v` occurs at most once in any case branch; the occurrence analyser
spots this and returns a OneOcc{ occ_n_br = 3 } for `v`. Then the code in
GHC.Core.Opt.Simplify.Utils.postInlineUnconditionally inlines `v` at its three
use sites, and discards the let-binding. That way, we avoid allocating `v` in
the A,B,C branches (though we still compute it of course), and branch D
doesn't involve <small thunk> at all. This sometimes makes a Really Big
Difference.
In (P2) we have shared the common RHS of A, B, in a join point `j`. We would
like to inline `v` in just the same way as in (P1). But the usual strategy
for let bindings is conservative and uses `andUDs` to combine usage from j's
RHS to its body; as if `j` was called on every code path (once, albeit). In
the case of (P2), we'll get ManyOccs for `v`. Important optimisation lost!
Solving this problem makes the Simplifier less fragile. For example,
the Simplifier might inline `j`, and convert (P2) into (P1)... or it might
not, depending in a perhaps-fragile way on the size of the join point.
I was motivated to implement this feature of the occurrence analyser
when trying to make optimisation join points simpler and more robust
(see e.g. #23627).
The occurrence analyser therefore has clever code that behaves just as
if you inlined `j` at all its call sites. Here is a tricky variant
to keep in mind:
Program (P3)
join j = case v of (a,b) -> a
in case f v of
A -> j
B -> j
C -> []
If you mentally inline `j` you'll see that `v` is used twice on the path
through A, so it should have ManyOcc. Bear this case in mind!
* We treat /non-recursive/ join points specially. Recursive join points are
treated like any other letrec, as before. Moreover, we only give this special
treatment to /pre-existing/ non-recursive join points, not the ones that we
discover for the first time in this sweep of the occurrence analyser.
* In occ_env, the new (occ_join_points :: IdEnv OccInfoEnv) maps
each in-scope non-recursive join point, such as `j` above, to
a "zeroed form" of its RHS's usage details. The "zeroed form"
* deletes ManyOccs
* maps a OneOcc to OneOcc{ occ_n_br = 0 }
In our example, occ_join_points will be extended with
[j :-> [v :-> OneOcc{occ_n_br=0}]]
See addJoinPoint.
* At an occurrence of a join point, we do everything as normal, but add in the
UsageDetails from the occ_join_points. See mkOneOcc.
* Crucially, at the NonRec binding of the join point, in `occAnalBind`, we use
`orUDs`, not `andUDs` to combine the usage from the RHS with the usage from
the body.
Here are the consequences
* Because of the perhaps-surprising OneOcc{occ_n_br=0} idea of the zeroed
form, the occ_n_br field of a OneOcc binder still counts the number of
/actual lexical occurrences/ of the variable. In Program P2, for example,
`v` will end up with OneOcc{occ_n_br=2}, not occ_n_br=3.
There are two lexical occurrences of `v`!
(NB: `orUDs` adds occ_n_br together, so occ_n_br=1 is impossible, too.)
* In the tricky (P3) we'll get an `andUDs` of
* OneOcc{occ_n_br=0} from the occurrences of `j`)
* OneOcc{occ_n_br=1} from the (f v)
These are `andUDs` together in `addOccInfo`, and hence
`v` gets ManyOccs, just as it should. Clever!
There are a couple of tricky wrinkles
(W1) Consider this example which shadows `j`:
join j = rhs in
in case x of { K j -> ..j..; ... }
Clearly when we come to the pattern `K j` we must drop the `j`
entry in occ_join_points.
This is done by `drop_shadowed_joins` in `addInScope`.
(W2) Consider this example which shadows `v`:
join j = ...v...
in case x of { K v -> ..j..; ... }
We can't make j's occurrences in the K alternative give rise to an
occurrence of `v` (via occ_join_points), because it'll just be deleted by
the `K v` pattern. Yikes. This is rare because shadowing is rare, but
it definitely can happen. Solution: when bringing `v` into scope at
the `K v` pattern, chuck out of occ_join_points any elements whose
UsageDetails mentions `v`. Instead, just `andUDs` all that usage in
right here.
This requires work in two places.
* In `preprocess_env`, we detect if the newly-bound variables intersect
the free vars of occ_join_points. (These free vars are conveniently
simply the domain of the OccInfoEnv for that join point.) If so,
we zap the entire occ_join_points.
* In `postprcess_uds`, we add the chucked-out join points to the
returned UsageDetails, with `andUDs`.
(W3) Consider this example, which shadows `j`, but this time in an argument
join j = rhs
in f (case x of { K j -> ...; ... })
We can zap the entire occ_join_points when looking at the argument,
because `j` can't posibly occur -- it's a join point! And the smaller
occ_join_points is, the better. Smaller to look up in mkOneOcc, and
more important, less looking-up when checking (W2).
This is done in setNonTailCtxt. It's important /not/ to do this for
join-point RHS's because of course `j` can occur there!
NB: this is just about efficiency: it is always safe /not/ to zap the
occ_join_points.
(W4) What if the join point binding has a stable unfolding, or RULES?
They are just alternative right-hand sides, and at each call site we
will use only one of them. So again, we can use `orUDs` to combine
usage info from all these alternatives RHSs.
Wrinkles (W1) and (W2) are very similar to Note [Binder swap] (BS3).