Comment améliorer les performances de ce calcul numérique dans Haskell?

Question

Je suis en cours de portage C implementation de David Blei d'allocation de Dirichlet latent à Haskell, et j'essaie de décider si laisser certains des éléments de bas niveau dans C. La fonction suivante est un exemple-c'est un approximation de la dérivée seconde de 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); 
} 

J'ai traduit cela en plus ou moins idiomatiques Haskell comme suit:

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) 

le problème est que quand je lance à la fois par Criterion, ma version Haskell est six ou sept fois plus lent r (Je compile avec -O2 sur GHC 6.12.1). Certaines fonctions similaires sont encore pire.

Je ne connais pratiquement rien aux performances de Haskell, et je ne m'intéresse pas vraiment à digging through Core ou à quoi que ce soit de ce genre, puisque je peux toujours appeler la poignée de fonctions C intensives en mathématiques par FFI. Mais je suis curieux de savoir s'il y a des fruits qui me manquent - une sorte d'extension ou de bibliothèque ou d'annotation que je pourrais utiliser pour accélérer ce truc numérique sans le rendre trop moche.

---

MISE À JOUR: Voici deux meilleures solutions, grâce à Don Stewart et Yitz. J'ai légèrement modifié la réponse de Yitz pour utiliser 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 

La performance des deux semble être presque exactement la même, avec l'un ou l'autre gagnant par un point de pourcentage ou deux selon les drapeaux du compilateur.

Comme camccann dit over at Reddit, la morale de l'histoire est "Pour de meilleurs résultats, utilisez Don Stewart comme générateur de code backend GHC." À part cette solution, le pari le plus sûr semble être juste de traduire les structures de contrôle C directement dans Haskell, bien que la fusion de boucles puisse donner des performances similaires dans un style plus idiomatique.

Je vais probablement utiliser l'approche Data.Vector dans mon code.

Answer 1

Utilisez les mêmes structures de contrôle et de données, ce qui donne:

{-# 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) 

Je n'ai pas votre testsuite, mais cela donne l'asm suivante:

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 

Ce qui a l'air ok. C'est le type de code que le backend -fllvm fait un bon travail. GCC déroule la boucle, et la seule façon de le faire est soit via le modèle Haskell ou le déroulement manuel. Vous pourriez considérer cela (une macro TH) si vous en faites beaucoup.

En fait, le backend LLVM GHC fait la boucle :-) déroulez

Enfin, si vous aimez vraiment la version originale Haskell, écrivez à l'aide stream fusion combinators, et GHC reconvertir en boucles. (Exercice pour le lecteur).

Answer 2

Avant le travail d'optimisation, je ne dirais pas que votre traduction originale est la façon la plus idiomatique d'exprimer dans Haskell ce que fait le code C.

Comment le processus d'optimisation ont procédé si nous avons commencé par ce qui suit à la place:

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