1 files changed, 19 insertions, 14 deletions
diff --git a/bench/Main.hs b/bench/Main.hs
index 5d2cb5a..af83ef7 100644
--- a/bench/Main.hs
+++ b/bench/Main.hs
@@ -18,24 +18,25 @@ import Data.Kind (Constraint)
 import GHC.Exts (withDict)
 
 import AST
+import AST.UnMonoid
 import Array
 import qualified CHAD (defaultConfig)
 import CHAD.Top
 import CHAD.Types
+import Compile
 import Data
 import Example
 import Example.GMM
 import Example.Types
-import Interpreter
 import Interpreter.Rep
 import Simplify
 
 
-gradCHAD :: KnownEnv env => CHADConfig -> SList Value env -> Double -> Ex env (TScal TF64) -> (Double, Rep (Tup (D2E env)))
-gradCHAD config input ctg term =
-  interpretOpen False input $
-    simplifyFix $
-      ELet ext (EConst ext STF64 ctg) $ chad' config knownEnv term
+gradCHAD :: KnownEnv env => CHADConfig -> Ex env (TScal TF64) -> IO (SList Value env -> IO (Double, Rep (Tup (D2E env))))
+gradCHAD config term =
+  compile knownEnv $
+    simplifyFix $ unMonoid $ simplifyFix $
+      ELet ext (EConst ext STF64 1.0) $ chad' config knownEnv term
 
 instance KnownTy t => NFData (Value t) where
   rnf = \(Value x) -> go (knownTy @t) x
@@ -115,12 +116,16 @@ accumConfig = chcSetAccum CHAD.defaultConfig
 
 main :: IO ()
 main = defaultMain
-  [env (return makeNeuralInputs) $ \inputs ->
-    bench "neural" (nf (\(inp, ctg) -> gradCHAD CHAD.defaultConfig inp ctg neural) (inputs, 1.0))
-  ,env (return makeNeuralInputs) $ \inputs ->
-    bench "neural-accum" (nf (\(inp, ctg) -> gradCHAD accumConfig inp ctg neural) (inputs, 1.0))
-  ,env (return makeGMMInputs) $ \inputs ->
-    bench "gmm" (nf (\(inp, ctg) -> gradCHAD CHAD.defaultConfig inp ctg (gmmObjective False)) (inputs, 1.0))
-  ,env (return makeGMMInputs) $ \inputs ->
-    bench "gmm-accum" (nf (\(inp, ctg) -> gradCHAD accumConfig inp ctg (gmmObjective False)) (inputs, 1.0))
+  [env (return makeNeuralInputs) $ \inputs -> bgroup "neural"
+    [env (gradCHAD CHAD.defaultConfig neural) $ \fun ->
+      bench "default" (nfAppIO fun inputs)
+    ,env (gradCHAD accumConfig neural) $ \fun ->
+      bench "accum" (nfAppIO fun inputs)
+    ]
+  ,env (return makeGMMInputs) $ \inputs -> bgroup "gmm"
+    [env (gradCHAD CHAD.defaultConfig (gmmObjective False)) $ \fun ->
+      bench "default" (nfAppIO fun inputs)
+    ,env (gradCHAD accumConfig (gmmObjective False)) $ \fun ->
+      bench "accum" (nfAppIO fun inputs)
+    ]
   ]