¿Cómo mejorar el rendimiento de este cálculo numérico en Haskell?

Question

Estoy en el medio de portar el original C implementation de David Blei de Latent Dirichlet Asignación a Haskell, y estoy tratando de decidir si dejar algunas de las cosas de bajo nivel en C. La siguiente función es un ejemplo: es un aproximación de la segunda derivada 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); 
} 

he traducido esto en mayor o menor idiomática Haskell de la siguiente manera:

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) 

el problema es que cuando corro tanto a través de Criterion, mi versión Haskell es seis o siete veces más lenta r (estoy compilando con -O2 en GHC 6.12.1). Algunas funciones similares son incluso peores.

No sé prácticamente nada sobre el rendimiento de Haskell, y no estoy muy interesado en digging through Core ni nada por el estilo, ya que siempre puedo llamar el puñado de funciones C intensivas en matemáticas a través de FFI.

Pero tengo curiosidad acerca de si hay poca fruta que me falta, alguna clase de extensión o biblioteca o anotación que podría usar para acelerar este material numérico sin hacerlo demasiado feo.

---

ACTUALIZACIÓN: Aquí hay dos soluciones mejores, gracias a Don Stewart y Yitz. He modificado ligeramente la respuesta de Yitz para usar 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 

El rendimiento de los dos parece ser casi exactamente el mismo, con uno u otro de ganar por un punto porcentual o dos, dependiendo de las opciones del compilador.

Como dijo camccannover at Reddit, la moraleja de la historia es "Para obtener los mejores resultados, utilice Don Stewart como su generador de código de fondo GHC". Salvo esa solución, la apuesta más segura parece ser simplemente trasladar las estructuras de control C directamente a Haskell, aunque la fusión de bucles puede ofrecer un rendimiento similar en un estilo más idiomático.

Probablemente terminaré utilizando el enfoque Data.Vector en mi código.

Answer 1

utilizar las mismas estructuras de control y de datos, obteniéndose:

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

no tengo el banco de pruebas, pero esto se obtiene la siguiente 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 

que se ve bien. Este es el tipo de código que el backend -fllvm hace un buen trabajo.

Sin embargo, GCC desenrolla el bucle, y la única manera de hacerlo es a través de Template Haskell o desenrollando manualmente. Puede considerar eso (una macro TH) si hace mucho de esto.

En realidad, el backend GHC LLVM hace desenrollar el bucle :-)

Por último, si realmente te gusta la versión original Haskell, lo escribe usando stream fusion combinators, y GHC lo convertirá de nuevo en bucles. (Ejercicio para el lector).

Answer 2

Antes del trabajo de optimización, no diría que su traducción original es la forma más idiomática de expresar en Haskell lo que hace el código C.

¿Cómo el proceso de optimización han procedido si empezamos con el siguiente lugar:

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