WorkWrap Idea: Unboxing higher-order function results based on strictness
Consider
{-# OPTIONS_GHC -O2 -fforce-recomp #-}
module Lib where
import System.Environment
f :: (Int -> Int) -> [Int] -> Int
f g [] = 0
f g (x:xs) = g x + f g xs
main = do
(n:_)<- map read <$> getArgs
print $ f (*n) [1..1000]
Currently, we get the following w/w split for f
:
$wf :: (Int -> Int) -> [Int] -> Int#
$wf g xs = ... case g x of I# n -> n +# ...
f g xs = I# ($wf g xs)
What I want is
$wf :: (Int -> Int#) -> [Int] -> Int#
$wf g' xs = ... case g' x :: Int# of n -> n +# ...
f g xs = I# ($wf (\n -> case g n of I# m -> m) xs)
Thus one fewer box to pass from g
to f
.
Note that polarity changes whether we can unbox args or results based on strictness. Here we may unbox the result of the higher-order function because we never use its box. Note also that we allocate an additional closure for the new lambda if f
doesn't inline, so this might not be ideal. Nevertheless, worth considering. With boxity analysis (#19871 (closed)), we should have all the demand analysis machinery we need to exploit this. E.g., f
puts a demand of LCL(!1L)
on g
, so we see that its result is used unboxed.
It's not unlike WW'ing for arity, where the wrapper has to allocate a closure for the squiggly lambda wrapper around the non-squiggly lambda call to g
. Here, we simply extend the idea to unbox the result of that call to g
.