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

meta name keywords

Question

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

Lo traduje a Haskell más o menos idiomático 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 ejecuto ambos a través de Criterion , mi versión de Haskell es seis o siete veces más lenta (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 profundizar en Core o algo por el estilo, ya que siempre puedo llamar a un puñado de funciones C intensivas en matemáticas a través de FFI.

Pero tengo curiosidad acerca de si hay alguna fruta que me falta, una especie 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 Data.Vector 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 ganar en un punto porcentual o dos dependiendo de los indicadores del compilador.

Como dijo camccann en Reddit , la moraleja de la historia es: "Para obtener los mejores resultados, usa Don Stewart como tu generador de código backend de 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é usando el enfoque Data.Vector en mi código.

Answer 1

Antes de la 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 habría sido el proceso de optimización si comenzáramos con lo siguiente en su 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

Answer 2

Use las mismas estructuras de control y datos, produciendo:

{-# 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 su suite de prueba, pero esto produce 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 -fllvm hace un buen trabajo.

Sin embargo, GCC desenrolla el bucle, y la única forma 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 de GHC LLVM desenrolla el ciclo :-)

Finalmente, si realmente te gusta la versión original de Haskell, escríbela usando los combinadores de fusión de flujo, y GHC la convertirá nuevamente en bucles. (Ejercicio para el lector).