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

Question

Я в процессе переноса оригинальной реализации C Дэвида Блея Скрытого распределения Дирихле в 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);
}

Я перевел это на более или менее идиоматический Хаскель следующим образом:

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)

Проблема в том, что когда я запускаю оба через Критерий , моя версия на Haskell работает в шесть или семь раз медленнее (я компилирую с -O2 на GHC 6.12.1). Некоторые похожие функции еще хуже.

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

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

---

ОБНОВЛЕНИЕ: Вот два лучших решения, благодаря Дону Стюарту и Ицу . Я немного изменил ответ Ица, чтобы использовать 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 в 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)

У меня нет вашего тестового набора, но это приводит к следующему 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 Combinator, и 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