Как повысить производительность этого численного вычисления в Haskell?

Question

Я нахожусь в середине портирования оригинала Дэвида Блей 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 в шесть или семь раз медленнее r (я компилирую с -O2 на GHC 6.12.1). Некоторые подобные функции еще хуже.

Я практически ничего не знаю о производительности Haskell, и меня не очень интересует digging through Core или что-то в этом роде, так как я всегда могу просто назвать кучу математических функций C через FFI.

Но мне любопытно, есть ли плохие фрукты, которые мне не хватает, - какое-то расширение или библиотека или аннотация, которые я мог бы использовать, чтобы ускорить эту цифровую штуку, не делая ее слишком уродливой.

---

UPDATE: Вот две лучшие решения, благодаря Don Stewart и 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, мораль истории - «Для достижения наилучших результатов используйте Дон Стюарт в качестве генератора кода GHC». Запрет на это решение, самая безопасная ставка, по-видимому, заключается в простом преобразовании структур управления C непосредственно в Haskell, хотя слияние фьюжн может придать аналогичную производительность в более идиоматическом стиле.

Я, вероятно, в конечном итоге воспользуюсь подходом Data.Vector в своем коде.

Answer 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) 

Я не ваш Тестов, но это приводит к следующему ассемблер:

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 разворачивает петлю, и единственный способ сделать это - либо с помощью шаблона Haskell, либо ручного разворота. Вы можете подумать, что (макрос TH), если вы делаете много этого.

На самом деле, бэкенд GHC LLVM делает раскатать петлю :-)

Наконец, если вы действительно, как в оригинальной версии Haskell, написать его с помощью stream fusion combinators, и GHC преобразует его обратно в петли. (Упражнение для читателя).

Answer 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