Wie verbessert man die Leistung dieser numerischen Berechnung in Haskell?

Question

Ich bin in der Mitte von David Blei Original C implementation Latent Dirichlet Allocation zu Haskell portieren, und ich versuche zu entscheiden, ob einige der Low-Level-Zeug in C verlassen. Die folgende Funktion ist ein Beispiel-es ist ein Annäherung der zweiten Ableitung von 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); 
} 

ich dies in mehr oder weniger idiomatische Haskell übersetzt haben, wie folgt:

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) 

Das Problem ist, dass, wenn ich laufen beide durch Criterion, meine Haskell-Version ist sechs oder sieben mal langsamer r (Ich kompiliere mit -O2 auf GHC 6.12.1). Einige ähnliche Funktionen sind noch schlimmer.

Ich weiß praktisch nichts über Haskell Leistung, und ich bin nicht sehr interessiert an digging through Core oder etwas ähnliches, da ich immer nur die Handvoll mathematisch intensive C-Funktionen über FFI aufrufen kann.

Aber ich bin neugierig, ob es tief hängende Früchte gibt, die ich vermisse - eine Art Erweiterung oder Bibliothek oder Annotation, mit denen ich diese Zahlen beschleunigen könnte, ohne sie zu hässlich zu machen.

---

UPDATE: Hier sind zwei bessere Lösungen, dank Don Stewart und Yitz. Ich habe Yitz 'Antwort leicht modifiziert, um Data.Vector zu verwenden.

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 

Die Leistung der beiden scheint fast genau die gleiche zu sein, mit dem einen oder anderen zu gewinnen um einen Prozentpunkt oder zwei auf den Compiler-Flags abhängig.

Wie camccann sagte over at Reddit, die Moral der Geschichte ist "Für beste Ergebnisse, verwenden Sie Don Stewart als Ihre GHC-Backend-Code-Generator." Abgesehen von dieser Lösung scheint die sicherste Wette nur darin zu bestehen, die C-Kontrollstrukturen direkt in Haskell zu übersetzen, obwohl die Schleifenfusion eine ähnliche Leistung in einem idiomatischen Stil liefern kann.

Ich werde wahrscheinlich am Ende mit der Data.Vector Ansatz in meinem Code.

Answer 1

Verwenden Sie die gleichen Steuer- und Datenstrukturen, wodurch man

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

Ich habe nicht Ihre Testsuite, aber dies ergibt die folgende 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 

Welche ok aussieht. Dies ist die Art von Code der -fllvm Backend macht einen guten Job.

GCC entrollt die Schleife jedoch, und der einzige Weg, dies zu tun, ist entweder über Vorlage Haskell oder manuelle Abrollung. Sie könnten dies (ein TH-Makro) betrachten, wenn Sie viel davon tun.

Eigentlich funktioniert das GHC LLVM Backend :-)

schließlich die Schleife entrollt, wenn Sie wirklich, wie die ursprüngliche Version Haskell, schreiben Sie es mit stream fusion combinators, und GHC wird es wieder in Schleifen konvertieren. (Übung für den Leser).

Answer 2

Bevor die Optimierung funktioniert, würde ich nicht sagen, dass Ihre ursprüngliche Übersetzung die idiomatische Art ist, in Haskell auszudrücken, was der C-Code macht.

Wie hätte der Optimierungsprozess fort, wenn wir mit den folgenden anstelle gestartet:

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