Come migliorare le prestazioni di questo calcolo numerico in Haskell?

Question

Sono nel mezzo del porting di David Blei originale C implementation di Latent Dirichlet Allocation a Haskell, e sto cercando di decidere se lasciare alcune delle cose di basso livello in C. La seguente funzione è un esempio: è un approssimazione della derivata seconda di 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); 
} 

ho tradotto questo in più o meno idiomatica Haskell come segue:

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) 

il problema è che quando ho eseguito sia attraverso Criterion, la mia versione Haskell è sei o sette volte lento r (Sto compilando con -O2 su GHC 6.12.1). Alcune funzioni simili sono anche peggiori.

Non conosco praticamente nulla delle prestazioni Haskell e non mi interessa molto lo digging through Core o qualcosa del genere, dal momento che posso sempre chiamare le poche funzioni C ad alta intensità di matematica tramite FFI.

Ma sono curioso di sapere se ci sono i frutti a bassa attaccatura che mi mancano: una specie di estensione o libreria o annotazione che potrei usare per accelerare questo materiale numerico senza renderlo troppo brutto.

---

UPDATE: Qui ci sono due soluzioni migliori, grazie a Don Stewart e Yitz. Ho modificato leggermente la risposta di Yitz per usare 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 

Le prestazioni dei due sembra essere quasi esattamente lo stesso, con uno o l'altro vincente di un punto percentuale o due a seconda delle opzioni del compilatore.

Come camccann come over at Reddit, la morale della storia è "Per ottenere i migliori risultati, utilizzare Don Stewart come generatore di codice di back-end GHC". Escludendo tale soluzione, la scommessa più sicura sembra essere quella di tradurre direttamente le strutture di controllo C in Haskell, sebbene la fusione in loop possa fornire prestazioni simili in uno stile più idiomatico.

Probabilmente finirò utilizzando l'approccio Data.Vector nel mio codice.

Answer 1

Utilizzare le stesse strutture di controllo e di dati, ottenendo:

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

non ho la vostra suite di test, ma questo produce il seguente 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 

Il che sembra ok. Questo è il tipo di codice che il backend -fllvm fa un buon lavoro.

GCC srotola il ciclo e l'unico modo per farlo è tramite Template Haskell o srotolamento manuale. Potresti considerare che (una macro TH) se fai molto di questo.

In realtà, il backend GHC LLVM fa srotolare il ciclo :-)

Infine, se vi piace la versione originale Haskell, scrive usando stream fusion combinators, e GHC sarà riconvertirlo in loop. (Esercizio per il lettore).

Answer 2

Prima del lavoro di ottimizzazione, non direi che la tua traduzione originale è il modo più idiomatico per esprimere in Haskell cosa sta facendo il codice C.

Come sarebbe il processo di ottimizzazione hanno proceduto se abbiamo iniziato con invece la seguente:

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