algorithm - 加速此功能的可能性是什么？

Question

我正在尝试加速以下功能:

{-# LANGUAGE BangPatterns #-}

import Data.Word
import Data.Bits
import Data.List (foldl1')
import System.Random
import qualified Data.List as L

data Tree a = AB (Tree a) (Tree a) | A (Tree a) | B (Tree a) | C !a
    deriving Show

merge :: Tree a -> Tree a -> Tree a
merge (C x) _               = C x
merge _ (C y)               = C y
merge (A ta) (A tb)         = A (merge ta tb)
merge (A ta) (B tb)         = AB ta tb
merge (A ta) (AB tb tc)     = AB (merge ta tb) tc
merge (B ta) (A tb)         = AB tb ta
merge (B ta) (B tb)         = B (merge ta tb)
merge (B ta) (AB tb tc)     = AB tb (merge ta tc)
merge (AB ta tb) (A tc)     = AB (merge ta tc) tb
merge (AB ta tb) (B tc)     = AB ta (merge tb tc)
merge (AB ta tb) (AB tc td) = AB (merge ta tc) (merge tb td)

为了强调其性能，我使用merge实现了排序:

fold ab a b c list = go list where
    go (AB a' b') = ab (go a') (go b')
    go (A a')     = a (go a')
    go (B b')     = b (go b')
    go (C x)      = c x

mergeAll :: [Tree a] -> Tree a
mergeAll = foldl1' merge

foldrBits :: (Word32 -> t -> t) -> t -> Word32 -> t
foldrBits cons nil word = go 32 word nil where
    go 0 w !r = r
    go l w !r = go (l-1) (shiftR w 1) (cons (w.&.1) r)

word32ToTree :: Word32 -> Tree Word32
word32ToTree w = foldrBits cons (C w) w where
    cons 0 t = A t
    cons 1 t = B t

toList = fold (++) id id (\ a -> [a])

sort = toList . mergeAll . map word32ToTree

main = do
    is <- mapM (const randomIO :: a -> IO Word32) [0..500000]
    print $ sum $ sort is

性能从一开始就非常好，比 Data.List 的 sort 慢 2.5 倍左右。不过，我没有做任何进一步加速的事情:内联几个函数，在许多地方添加注释，UNPACK on C !a。有什么办法可以加快这个功能吗？

Answer 1

您肯定分配了太多 thunk。我将展示如何分析代码:

merge (A ta) (A tb)         = A (merge ta tb)

在这里，您为构造函数 A 分配了一个参数，它是一个 thunk。你能说什么时候 merge ta tb block 会被强制执行吗？可能只有在最后，使用生成的树时。尝试为每个 Tree 构造函数的每个参数添加一个 bang 以确保它是 spine-strict:

data Tree a = AB !(Tree a) !(Tree a) | A !(Tree a) | B !(Tree a) | C !a

下一个例子:

go l w !r = go (l-1) (shiftR w 1) (cons (w.&.1) r)

在这里，您正在为 l-1、shifrR w 1 和 cons (w.&.1) r 分配一个 thunk。当比较 l 与 o 时，第一个将在下一次迭代中强制执行，第二个将在下一次迭代中强制执行 3d thunk 时强制执行 (w 在这里使用)，并且由于 r 上的爆炸，第三个 thunk 被强制用于下一次迭代。所以可能这个特定的条款没问题。