Add peeling/unrolling support
Motivation
Currently we can't inline recursive functions. It would be nice to add a way to automatically peel and unroll recursive functions (just like loops in imperative languages).
Suppose we have:
Rec { f x = body -- Where 'body' mentions 'f'
}
We can peel off one iteration like this:
Rec { $u_f x = body[$u_f/f] }
NonRec { f x = body[$u_f/f] }
where $u_f
is the unrolled worker of f
(not unrolled yet in this case) and f
is the peeled off wrapper. The wrapper is non recursive and may be inlined.
We could do this process again to peel off several iterations:
Rec { $u_f x = body[$u_f/f] }
NonRec { f x = body[f0/f]
, f0 x = body[f1/f]
, f1 x = body[$u_f/f]
}
Similarly we could unroll in the worker too:
Rec { $u_f x = body[$u_f0/f] -- loop-breaker
, $u_f0 x = body[$u_f1/f]
, $u_f1 x = body[$u_f2/f]
, $u_f2 x = body[$u_f/f]
}
NonRec { f x = body[f0/f]
, f0 x = body[f1/f]
, f1 x = body[$u_f/f]
}
fN
and $u_fN
bindings will probably become join points (or be inlined if there is only one recursive call in body
), be specialised, etc.
The body
of the worker can be specialised before being unrolled: we know all the call-sites so we can remove dead-code, perform worker-wrapper, etc.
We can also have one worker per call-site to specialise each even more.
Example
Example from GHC.Utils.Ppr
that could benefit from this:
-- recursive
reduceDoc :: Doc -> RDoc
reduceDoc (Beside !p !g q) = beside p g $! reduceDoc q
reduceDoc (Above !p !g q) = above p g $! reduceDoc q
reduceDoc p = p
-- recursive
beside :: Doc -> Bool -> RDoc -> RDoc
beside NoDoc _ _ = NoDoc
beside (p1 `Union` p2) g q = beside p1 g q `union_` beside p2 g q
beside Empty _ q = q
beside (Nest k p) g q = nest_ k $! beside p g q
beside p@(Beside p1 g1 q1) g2 q2
| g1 == g2 = beside p1 g1 $! beside q1 g2 q2
| otherwise = beside (reduceDoc p) g2 q2
beside p@(Above{}) g q = let !d = reduceDoc p in beside d g q
beside (NilAbove p) g q = nilAbove_ $! beside p g q
beside (TextBeside s sl p) g q = textBeside_ s sl rest
where
rest = case p of
Empty -> nilBeside g q
_ -> beside p g q
In calls to reduceDoc p
in beside
, p
is already scrutinised so it would be beneficial to inline reduceDoc
's peeled off wrapper.
Another example from the same module:
mkNest :: Int -> Doc -> Doc
mkNest !k (Nest k1 p) = mkNest (k + k1) p
mkNest _ NoDoc = NoDoc
mkNest _ Empty = Empty
mkNest 0 p = p
mkNest k p = Nest k p
-- smart constructor for Nest
nest :: Int -> Doc -> Doc
nest k p = mkNest k (reduceDoc p)
mkNest
won't inline because it is recursive. But it would be reasonable to inline mkNest
and reduceDoc
wrappers in nest
so that case-of-case transformation could transform nest
into something like:
nest :: Int -> Doc -> Doc
nest !k = \case
Beside !p !g q -> case beside p g $! $u_reduceDoc q of ...
Above !p !g q -> case above p g $! $u_reduceDoc q of ...
Nest k1 p -> $u_mkNest (k+k1) p
NoDoc -> NoDoc
Empty -> Empty
p | k == 0 -> p
p -> Nest k p
Proposal
Add the following pragma:
{-# UNROLL [u] [[SPECIALISE] PEEL p] identifier #-}
Where:
- the optional
u >= 1
(default to 1) is the number of unrolled iterations in the worker -
p >= 1
is the number of peeled off iterations in the wrapper - SPECIALISE can be used to require one worker per recursive call in the wrapper (to specialise each worker for each call-site)
Note that INLINE pragmas can be applied to the wrapper.
Example:
{-# UNROLL PEEL 1 reduceDoc #-}
{-# INLINE reduceDoc #-}