Custom Fusion Optimizations in Haskell
An important component to improving the performance of lists in Haskell is the idea of fusion optimizations. These are special types of optimizations that eliminate intermediate structures such as lists to reduce the total amount of work done in a program and also allow for other optimizations such as unboxing.
Consider the following Haskell expression:
sum [1 .. 1000000]This expression evaluates to the sum of all integers between 1 and 1,000,000
(inclusive). Due to lazy evaluation, Haskell, with optimizations disabled, does
not create a list of 1 million elements all at once which saves on memory, but
does not save any time since each cons cell and Int box are still created.
In situations like these, performance for Haskell becomes very poor compared to the “fast” languages like C/C++/Rust/etc since for each addition, every number is boxed, put into a cons cell, then extracted from a cons cell, and then unboxed. You may notice that boxing the value and placing it into a cons cell is unnecessary work since it is immediately extracted, and you would be correct; the core idea of fusion is to have these types of actions cancel each other out.
When turning on optimizations in GHC, the above expression is inlined to resemble a producer/consumer relationship which is then fused together via a rewrite rule to yield an expression that is easily unboxable and has nearly identical performance to C.
Rewrite Rules
In Haskell, Fusion is performed by applying rewrite rules to the source code at compile time. One can think of them as a fancy find-and-replace feature with some extra features.
The syntax for defining rules is <rulename> <lhs> = <rhs> wrapped inside a
RULES pragma. Here is an example of a rewrite rule that replaces all instances
of length [] with 0:
{-# RULES "length[]/0" length [] = 0 #-}GHC does not check for equivalence before and after the rewrite, so it is
possible to have a rule like length [] = 999. You can verify this yourself by
running print (length []) with optimizations enabled and see that it will
print 999 instead of 0. Rewrite rules are type-checked, however, so it is not
possible to say length [] = False without getting a compiler error.
foldr/build Rule
Alongside user-specified rewrite rules, GHC also contains a large amount of
built in rules for fusion and many other things. One of the most important rules
for lists is the foldr/build rule:
{-# RULES "foldr/build" forall k z g. foldr k z (build g) = g k z #-}At a high level, foldr/build says, “if you give me a function which produces
a list and a function that folds the list into a single value, I will give you
a function that does both at the same time without the intermediate list.”
To understand exactly how this works, lets look at how foldr and build
work.
The foldr function whose type signature is provided below takes an initial
accumulator, a function which updates the accumulator from an element of a
list, the input list, and returns the final accumulator after going through all
elements of the list.
-- Specialized to lists
foldr :: (a -> b -> b) -> b -> [a] -> b
foldr _ z [] = z
foldr f z (x:xs) = f x (foldr f z xs)The build function is a bit more complicated. It takes in a polymorphic
function as input which itself takes in a cons-like function, a nil-like value,
and returns l. This function is then applied to (:) and [] where (:) is
the cons-like function, and [] is the nil-like value. In a more technical sense,
it takes in a Church-encoded list as input and returns a regular list as output.
build :: (forall l . (a -> l -> l) -> l -> l) -> [a]
build builder = builder (:) []To get a better intuition, lets look at an example builder function which builds
the list [1, 2, 3]:
builder = \c n -> c 1 (c 2 (c 3 n)) This is very similar to building a list out of (:) and [], only c and n
could be anything. For example, lets use (+) as the cons-like function and 0
as the nil-like value:
>>> builder (+) 0
(\c n -> c 1 (c 2 (c 3 n))) (+) 0
(+) 1 ((+) 2 ((+) 3 0))
(1 + (2 + (3 + 0)))
6By doing this, we’ve effectively computed the sum of the list without
actually constructing the list at all; we’ve performed fusion. Revisiting the
foldlr/build rule from earlier:
{-# RULES "foldr/build" forall k z g. foldr k z (build g) = g k z #-}We see that we can do the same for any pair of cons-like functions and nil-like
values by using foldr. k is the cons-like function, z is the nil-like
value, and g is the builder function.
Functions like sum or [a .. b] are expressed with foldr and build so
that complex expressions such as map (+1) . take 5 . filter even are
fused together during compilation.
Designing a Custom Fusion Rule
GHC typically focuses on lists for fusion optimizations, but we saw earlier that rewrite rules can be user-specified, which means that the user can create their own fusion optimizations for their own types.
Tree
Let’s define a basic Tree structure typically used for binary searches.
Though for simplicity, we will not enforce any convention for how values
are placed into the tree.
data Tree a = Node a (Tree a) (Tree a) | End deriving ShowbuildTree - Building a Tree
To implement fusion, we’ll need to implement the corresponding build and fold functions. The build function is the simplest, as it simply converts the Church-encoded tree into a regular tree:
{-# INLINE[0] buildTree #-}
buildTree :: (forall t . (a -> t -> t -> t) -> t -> t) -> Tree a
buildTree builder = builder Node EndThere is an INLINE pragma for buildTree, which says, “do not inline
buildTree until the final phase of optimization.” This will give the fusion
rule, when we write it, to have a chance to run. If buildTree is inlined
before that, then we may miss the opportunity for fusion.
Let’s also implement another function, treeFromList that we can use to test
fusion on later and implement it in terms of buildTree. Note the use of
foldl', which internally is implemented as foldr which will allow for
fusion to occur with the input list.
{-# INLINE treeFromList #-}
treeFromList :: [a] -> Tree a
treeFromList xs = buildTree $ \node end ->
foldl' (\t a -> node a t end) end xstreeFromList also ignores the input values when inserting to the tree for
the purposes of simplicity. We add an INLINE pragma to treeFromList so that
it expands to buildTree as quickly as possible during compilation.
tree - Folding a Tree
Now we need to implement a function to fold the tree. For lists this is the
foldr function, for Maybe this is the maybe function, and for Either
this is the either function. To match the convention, we will call our
function tree. This function will simply substitute the node and end
functions in place of the Node and End constructors.
{-# INLINE[0] tree #-}
tree :: (a -> t -> t -> t) -> t -> Tree a -> t
tree node end End = end
tree node end (Node x lhs rhs) =
node x (tree node end lhs) (tree node end rhs)Tree Rewrite Rule
Before writing a rewrite rule, lets first see if the compiler can do better. Below is a function which we will use to test fusion that first creates a list of the numbers between 1 and 1,000,000 inclusive, converts it into a tree, and then returns the sum of the tree. This is two intermediate structures in total, and we would like to completely eliminate both of them. We mark it as inline so that GHC will focus on optimizing it instead of where to inline it.
{-# INLINE treeSum #-}
treeSum :: Num a => Tree a -> a
treeSum = tree (\x lhs rhs -> x + lhs + rhs) 0
{-# NOINLINE treeFusionTest #-}
treeFusionTest :: Int
treeFusionTest = treeSum (treeFromList [1 .. 1000000])When compiling with optimizations, we get the following core output:
-- RHS size: {terms: 3, types: 1, coercions: 0, joins: 0/0}
treeFusionTest [InlPrag=NOINLINE] :: Int
[GblId]
treeFusionTest = go1_r3U3 lvl23_r3U4 (Main.End @Int)
Rec {
-- RHS size: {terms: 19, types: 15, coercions: 0, joins: 0/0}
go1_r3U3 :: [Int] -> Tree Int -> Int
[GblId[StrictWorker([!, !])],
Arity=2,
Str=<1L><1L>,
Unf=OtherCon []]
go1_r3U3
= \ (ds_a3M3 :: [Int]) (eta_B0 [OS=OneShot] :: Tree Int) ->
case ds_a3M3 of {
[] -> tree @Int @Int arg_r3U2 Main.$fFoldableTree4 eta_B0;
: y_a3M6 ys_a3M7 ->
case eta_B0 of z_a2P5 { __DEFAULT ->
go1_r3U3 ys_a3M7 (Main.Node @Int y_a3M6 z_a2P5 (Main.End @Int))
}
}
end Rec }We see that GHC was able to do some inlining, but not
do much else after that. go1_r3U3 is still boxed, for example, which means
that it has to unbox all its inputs and box its output every time it is called,
and uses both intermediate structures we were trying to eliminate.
Knowing this, lets write our own fusion rule so that GHC can do a better job.
Following the template of the foldr/build rule, we need to find an expression
such that a tree that is built is immediately folded int a value:
tree node end (buildTree builder)In the above expression, buildTree builder builds the tree with Node and
End, but they are then immediately replaced with the node and end
functions. So the fusion here would be to just apply node and end to the
builder function directly:
tree node end (buildTree builder) = builder node endNow, we can write the fusion rule:
{-# RULES
"tree/buildTree" forall node end
(builder :: forall t . (a -> t -> t -> t) -> t -> t)
. tree node end (buildTree builder) = builder node end
#-}(Note that we have to explicitly specify the type of builder due to it being
a Rank-2 type)
Fusion Results
Now let’s look at the core output for treeFusionTest with the fusion rule:
-- RHS size: {terms: 7, types: 1, coercions: 0, joins: 0/0}
treeFusionTest [InlPrag=NOINLINE] :: Int
[GblId]
treeFusionTest
= case $wgo3_r3UT 1# 0# of ww_s3Tl { __DEFAULT ->
GHC.Types.I# ww_s3Tl
}
Rec {
-- RHS size: {terms: 16, types: 3, coercions: 0, joins: 0/0}
$wgo3_r3UT :: Int# -> Int# -> Int#
[GblId, Arity=2, Str=<1L><L>, Unf=OtherCon []]
$wgo3_r3UT
= \ (x_s3SV :: Int#) (ww_s3SY :: Int#) ->
case x_s3SV of wild_X1E {
__DEFAULT -> $wgo3_r3UT (+# wild_X1E 1#) (+# wild_X1E ww_s3SY);
100# -> +# 100# ww_s3SY
}
end Rec }Now, we see the apart from the wrapper, due to the use of NOINLINE, the
entire function has been reduced to primitive operations and, apart from the
wrapper, does not have any boxing at all. This code will perform much better
than without the fusion rule.
This function also does not have any heap or stack frame allocations, so it will have the same performance as a C for loop which adds all the values to a single variable. Also note that a true one-to-one mapping of the Haskell algorithm in C would actually lead to Haskell performing better than C since C cannot perform fusion optimizations (without abusing macros).
Conclusion
Fusion is a critically important optimization for removing intermediate structures and enabling further optimizations to occur. When used in conjunction with Haskell’s other optimizations, it can produce code with the same performance as C. Custom fusion operations are also only possible due to the extremely underrated feature of rewrite rules, which themselves are not just limited to fusion.