Optimise by backpropagating scalar computation in C

1 files changed, 44 insertions, 13 deletions

diff --git a/bench/Main.hs b/bench/Main.hs
index 99c3f1d..1174a3a 100644
--- a/bench/Main.hs
+++ b/bench/Main.hs
@@ -1,8 +1,14 @@
 {-# LANGUAGE TypeApplications #-}
 module Main where
 
+import Control.DeepSeq
+import Control.Exception (evaluate)
+import Control.Monad (forM_)
 import Criterion
 import Criterion.Main
+import qualified System.Clock as Clock
+import System.Environment (getArgs)
+import System.Mem (performGC)
 
 import qualified Numeric.AD as AD
 
@@ -11,17 +17,42 @@ import Numeric.ADDual.Examples
 
 
 main :: IO ()
-main = defaultMain
-  [env (pure (makeNeuralInput 100)) $ \input ->
-    bgroup "neural-100"
-      [bench "dual" $ nf (\inp -> ADD.gradient' @Double fneural inp 1.0) input
-      ,bench "ad" $ nf (\inp -> AD.grad fneural inp) input]
-  ,env (pure (makeNeuralInput 500)) $ \input ->
-    bgroup "neural-500"
-      [bench "dual" $ nf (\inp -> ADD.gradient' @Double fneural inp 1.0) input
-      ,bench "ad" $ nf (\inp -> AD.grad fneural inp) input]
-  ,env (pure (makeNeuralInput 2000)) $ \input ->
-    bgroup "neural-2000"
-      [bench "dual" $ nf (\inp -> ADD.gradient' @Double fneural inp 1.0) input
-      ,bench "ad" $ nf (\inp -> AD.grad fneural inp) input]
+main = do
+  args <- getArgs
+  case args of
+    ["--neural-graph"] -> mainNeuralGraph
+    _ -> mainCriterion
+
+mainCriterion :: IO ()
+mainCriterion = defaultMain
+  [benchNeural 100
+  ,benchNeural 180  -- rather stably 2 GCs
+  ,benchNeural 500
+  ,benchNeural 2000
   ]
+  where
+    benchNeural :: Int -> Benchmark
+    benchNeural n =
+      env (pure (makeNeuralInput n)) $ \input ->
+        bgroup ("neural-" ++ show n)
+          [bench "dual" $ nf (\inp -> ADD.gradient' fneural inp 1.0) input
+          ,bench "ad" $ nf (\inp -> AD.grad fneural inp) input]
+
+mainNeuralGraph :: IO ()
+mainNeuralGraph = do
+  forM_ [10, 20 .. 300] $ \n -> do
+    let input = makeNeuralInput n
+    _ <- evaluate (force input)
+
+    performGC
+    t1 <- Clock.getTime Clock.Monotonic
+    _ <- evaluate $ force (ADD.gradient' fneural input 1.0)
+    t2 <- Clock.getTime Clock.Monotonic
+
+    performGC
+    t3 <- Clock.getTime Clock.Monotonic
+    _ <- evaluate $ force (AD.grad fneural input)
+    t4 <- Clock.getTime Clock.Monotonic
+
+    let diff a b = fromIntegral (Clock.toNanoSecs (Clock.diffTimeSpec a b)) / 1e9 :: Double
+    putStrLn $ show n ++ " " ++ show (diff t1 t2) ++ " " ++ show (diff t3 t4)