Implement Nested CPR

This ticket tracks the design and implementation of a Nested CPR, e.g. extending CPR analysis to unbox constructed products nested inside constructed products.

A list of tickets for motivation:

• CPR IO/State computations #1600 (closed)
• Various places in GHC/base where we currently unbox IO manually, like GHC.Utils.Encoding or UniqSM
• A simple rapid-termination analysis (higher-order exprIsCheap/exprOkForSpeculation) #8655
• Effective sum CPR needs to be nested and hence the terminaion analysis machinery
• Many heuristics of our current CPR analysis (i.e. Note [CPR for binders that will be unboxed]) are an artifact of not being able to express and propagate nested CPRs in case scrutinees to field binders.

added CPR/boxity analysis needs triage labels

mentioned in merge request !1866 (closed)

changed the description

marked this issue as related to #16335

added Ttask label and removed needs triage label

mentioned in issue #18154

changed the description

changed the description

@lexi.lambda brought up an interesting extension to Nested CPR, see https://mail.haskell.org/pipermail/ghc-devs/2020-December/019406.html.

If we have

f :: Int -> (Int, (Int -> Int))
f a = (a+1, \b -> a + b)

then instead of just unboxing the first component, we could also w/w the lambda in the second component:

f (I# a) = case $wf a of (# a', $wg #) -> ( I# a', (\(I# b) -> case $wg b of r' -> I# r') #)

$wf :: Int# -> (# Int#, (Int# -> Int#) #)
$wf a = (# a +# 1#, \b -> a +# b #)

But that potentially allocates an additional lambda in the wrapper f that might not cancel away in higher order scenarios (e.g. when we call map g where g is the returned lambda). Not sure if we have the analysis info available that could make this process conservative. It's not simple demand info as I was hoping for in the reply to the above mail. But maybe we just shouldn't care for the higher-order scenario, not sure.

Another insight: I think Alexis was on point when suggesting to make Nested CPR "higher-order": In order not to allocate an additional lambda for call sites like

f :: Int -> (Int, (Int -> Int))
f a = (a+1, \b -> a + b)

main = ... case f (I# m) of (_, g) -> map g blah ...

==> { CPR nestedly, for the lambda }

f (I# a) = case $wf a of (# a', $wg #) -> ( I# a', (\(I# b) -> case $wg b of r' -> I# r') #)

$wf :: Int# -> (# Int#, (Int# -> Int#) #)
$wf a = (# a +# 1#, \b -> a +# b #)

main = ... case $wf m of (# _, $wg #) -> map (\(I# b) -> $wg b) blah ....

We have to see that map does not unpack the Int returned by its higher-order argument. That's exactly the boxity information we encode in a product demand! In the example above, main puts demand C1(U) on g. NB: It is not C1(P(U)), e.g. a product use on g's result, which would allow us to unbox the result of g without remorse. So we wouldn't w/w for the lambda in the example above, because the demand says "we use the box", ergo the lambda (which is really the wrapper for the lambda in f) won't cancel away.

Here's another example where it could cancel away:

data IntList = Cons {-# UNPACK #-} !Int IntList | Nil

mapInt :: (Int -> Int) -> IntList -> IntList
mapInt f Nil = Nil
mapInt f (Cons x xs)  = Cons (case (f (I# x)) of I# r -> r) (mapInt f xs)

f :: Int -> (Int, (Int -> Int))
f a = (a+1, \b -> a + b)

main = ... case f (I# m) of (_, g) -> mapInt g blah ...

==> { CPR nestedly, for the lambda, also transform mapInt }

mapInt :: (Int -> Int) -> IntList -> IntList
mapInt f = $wmapInt (\i -> case f (I# i) of I# r -> r)

$wmapInt :: (Int# -> Int#) -> IntList -> IntList
$wmapInt f Nil = Nil
$wmapInt f (Cons x xs) = Cons (f x) (mapInt f xs)

f (I# a) = case $wf a of (# a', $wg #) -> ( I# a', (\(I# b) -> case $wg b of r' -> I# r') #)

$wf :: Int# -> (# Int#, (Int# -> Int#) #)
$wf a = (# a +# 1#, \b -> a +# b #)

main = ... case $wf m of (# _, $wg #) -> $wmapInt $wg blah ....

Note that we had to transform mapInt for its demand on f (C1(P(U))). It's yet unclear to me if that's always an improvement, but I suppose it is.

mentioned in commit 9873fa2a

mentioned in commit 2682896b

mentioned in commit 4c15c483

mentioned in commit 3a0802dc

mentioned in commit f940f873

mentioned in commit 271c6bf2

mentioned in commit 0b64a0ac

mentioned in commit 3fec1c17

Having another look at #18174 (comment 318048), I think we could also just let bind the lambda before DmdAnal/CprAnal. Then we only have to realise that gs wrapper can inline after WW, thus

f :: Int -> (Int, (Int -> Int))
f a = (a+1, let g b = a + b in g)

==> { Nested CPR + WW }

$wf :: Int# -> (# Int#, (Int# -> Int#) #)
$wf a = (# a +# 1#, let $wg b = a +# b #)

f (I# a) = case $wf a of (# a', $wg #) ->
  (I# a', let g (I# b) = I# ($wg b) in g)

Currently, Nested CPR only considers things of arity 0 for WW, but this shows that we also want to transform fields of arity > 0, like g. Although it's not so clear if that's always a win, because we allocate $wg's closure in addition to g's. But the difference is merely in the closure for the wrapper, which will only close over $wg. Maybe the cases where the wrapper doesn't inline are rare.

This reminds me a bit of KaCC-sytyle WW for higher-order function arguments, only "the other way round", e.g. for results rather than arguments.

mentioned in commit bc3fe68c