如何在 Haskell 中提高这种数值计算的性能？答案

【问题标题】：How to improve performance of this numerical computation in Haskell?如何在 Haskell 中提高这种数值计算的性能？
【发布时间】：2011-02-28 01:40:01
【问题描述】：

我正在将 David Blei 的 Latent Dirichlet Allocation 的原始 C implementation 移植到 Haskell，我正在尝试决定是否将一些低级的东西留在 C 中。以下函数是一个示例——它是lgamma 的二阶导数的近似值：

double trigamma(double x)
{
    double p;
    int i;

    x=x+6;
    p=1/(x*x);
    p=(((((0.075757575757576*p-0.033333333333333)*p+0.0238095238095238)
         *p-0.033333333333333)*p+0.166666666666667)*p+1)/x+0.5*p;
    for (i=0; i<6 ;i++)
    {
        x=x-1;
        p=1/(x*x)+p;
    }
    return(p);
}

我已将其翻译成或多或少惯用的 Haskell，如下所示：

trigamma :: Double -> Double
trigamma x = snd $ last $ take 7 $ iterate next (x' - 1, p')
  where
    x' = x + 6
    p  = 1 / x' ^ 2
    p' = p / 2 + c / x'
    c  = foldr1 (\a b -> (a + b * p)) [1, 1/6, -1/30, 1/42, -1/30, 5/66]
    next (x, p) = (x - 1, 1 / x ^ 2 + p)

问题是，当我同时通过Criterion 运行时，我的Haskell 版本慢了六到七倍（我在GHC 6.12.1 上使用-O2 进行编译）。一些类似的功能就更差了。

我对 Haskell 性能几乎一无所知，而且我对 digging through Core 或类似的东西也不是很感兴趣，因为我总是可以通过 FFI 调用少数数学密集型 C 函数。

但我很好奇我是否缺少一些容易实现的成果——某种扩展、库或注释，我可以使用它们来加速这些数字内容，而不会使它变得太丑陋。

更新：这里有两个更好的解决方案，感谢Don Stewart 和Yitz。我稍微修改了 Yitz 的答案以使用Data.Vector。

invSq x = 1 / (x * x)
computeP x = (((((5/66*p-1/30)*p+1/42)*p-1/30)*p+1/6)*p+1)/x+0.5*p
  where p = invSq x

trigamma_d :: Double -> Double
trigamma_d x = go 0 (x + 5) $ computeP $ x + 6
  where
    go :: Int -> Double -> Double -> Double
    go !i !x !p
        | i >= 6    = p
        | otherwise = go (i+1) (x-1) (1 / (x*x) + p)

trigamma_y :: Double -> Double
trigamma_y x = V.foldl' (+) (computeP $ x + 6) $ V.map invSq $ V.enumFromN x 6

两者的性能似乎几乎完全相同，根据编译器标志，一个或另一个获胜一个或两个百分点。

正如camccann 所说的over at Reddit，故事的寓意是“为了获得最佳结果，请使用 Don Stewart 作为您的 GHC 后端代码生成器。”除了这种解决方案，最安全的选择似乎是将 C 控制结构直接转换为 Haskell，尽管循环融合可以以更惯用的风格提供类似的性能。

我可能最终会在我的代码中使用Data.Vector 方法。

【问题讨论】：

C 程序使用循环，而在 Haskell 中您使用的是堆分配列表。他们不会有同样的表现。最好的做法是直接将控件和数据结构翻译成 Haskell，以保持相同的性能。
嗨，特拉维斯！完成后你会发布你的代码吗？我发现我可以根据 C 代码理解你的 Haskell.. 也许我可以通过这种方式学习 Haskell..
您应该查看 FastInvSqrt 代码。
@朱印：是的，当然！完成后我会在这里发布一个链接。希望对你有用。
@DeadMG：感谢您的建议，但我不确定这是否会有所帮助，因为我们在这里不取平方根。我错过了什么吗？

标签： c performance math haskell

【解决方案1】：

使用相同的控制和数据结构，产生：

{-# LANGUAGE BangPatterns #-}
{-# OPTIONS_GHC -fvia-C -optc-O3 -fexcess-precision -optc-march=native #-}

{-# INLINE trigamma #-}
trigamma :: Double -> Double
trigamma x = go 0 (x' - 1) p'
    where
        x' = x + 6
        p  = 1 / (x' * x')

        p' =(((((0.075757575757576*p-0.033333333333333)*p+0.0238095238095238)
                  *p-0.033333333333333)*p+0.166666666666667)*p+1)/x'+0.5*p

        go :: Int -> Double -> Double -> Double
        go !i !x !p
            | i >= 6    = p
            | otherwise = go (i+1) (x-1) (1 / (x*x) + p)

我没有你的测试套件，但这会产生以下 asm：

A_zdwgo_info:
        cmpq    $5, %r14
        jg      .L3
        movsd   .LC0(%rip), %xmm7
        movapd  %xmm5, %xmm8
        movapd  %xmm7, %xmm9
        mulsd   %xmm5, %xmm8
        leaq    1(%r14), %r14
        divsd   %xmm8, %xmm9
        subsd   %xmm7, %xmm5
        addsd   %xmm9, %xmm6
        jmp     A_zdwgo_info

看起来不错。这是 -fllvm 后端做得很好的那种代码。

GCC 会展开循环，唯一的方法是通过 Template Haskell 或手动展开。如果做很多这样的事情，你可能会考虑（一个 TH 宏）。

实际上，GHC LLVM 后端确实展开了循环 :-)

最后，如果你真的喜欢原来的 Haskell 版本，使用stream fusion combinators, 编写它，GHC 会将它转换回循环。（读者练习）。

【讨论】：

谢谢，唐——这太棒了。在我的测试设置中，您的版本实际上以某种方式击败了 C 版本（略微）。不过，作为记录，第一行应为trigamma x = go 0 (x' - 1) p'，p 和p' 定义中的x 实例应替换为x'。
只是出于兴趣，您是否使用遗传算法来获得那些编译选项？

【解决方案2】：

在优化工作之前，我不会说您的原始翻译是在 Haskell 中表达 C 代码正在做什么的最惯用的方式。

如果我们从以下内容开始，优化过程将如何进行：

trigamma :: Double -> Double
trigamma x = foldl' (+) p' . map invSq . take 6 . iterate (+ 1) $ x
where
  invSq y = 1 / (y * y)
  x' = x + 6
  p  = invSq x'
  p' =(((((0.075757575757576*p-0.033333333333333)*p+0.0238095238095238)
              *p-0.033333333333333)*p+0.166666666666667)*p+1)/x'+0.5*p

【讨论】：