Benchmark

4 files changed, 74 insertions, 21 deletions

diff --git a/src/Array.hs b/src/Array.hs
index 8507544..ef9bb8d 100644
--- a/src/Array.hs
+++ b/src/Array.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE DeriveTraversable #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE KindSignatures #-}
@@ -6,10 +7,12 @@
 {-# LANGUAGE TupleSections #-}
 module Array where
 
+import Control.DeepSeq
 import Control.Monad.Trans.State.Strict
 import Data.Foldable (traverse_)
 import Data.Vector (Vector)
 import qualified Data.Vector as V
+import GHC.Generics (Generic)
 
 import Data
 
@@ -20,12 +23,20 @@ data Shape n where
 deriving instance Show (Shape n)
 deriving instance Eq (Shape n)
 
+instance NFData (Shape n) where
+  rnf ShNil = ()
+  rnf (sh `ShCons` n) = rnf n `seq` rnf sh
+
 data Index n where
   IxNil :: Index Z
   IxCons :: Index n -> Int -> Index (S n)
 deriving instance Show (Index n)
 deriving instance Eq (Index n)
 
+instance NFData (Index n) where
+  rnf IxNil = ()
+  rnf (sh `IxCons` n) = rnf n `seq` rnf sh
+
 shapeSize :: Shape n -> Int
 shapeSize ShNil = 1
 shapeSize (ShCons sh n) = shapeSize sh * n
@@ -51,7 +62,8 @@ enumShape sh = map (fromLinearIndex sh) [0 .. shapeSize sh - 1]
 
 -- | TODO: this Vector is a boxed vector, which is horrendously inefficient.
 data Array (n :: Nat) t = Array (Shape n) (Vector t)
-  deriving (Show, Functor, Foldable, Traversable)
+  deriving (Show, Functor, Foldable, Traversable, Generic)
+instance NFData t => NFData (Array n t)
 
 arrayShape :: Array n t -> Shape n
 arrayShape (Array sh _) = sh
diff --git a/src/CHAD.hs b/src/CHAD.hs
index 6b0627d..ffbdcac 100644
--- a/src/CHAD.hs
+++ b/src/CHAD.hs
@@ -689,22 +689,20 @@ retConcat descr (SCons (Ret (b :: Bindings _ _ shbinds1) (subtape :: Subenv _ ta
 
 freezeRet :: Descr env sto
           -> Ret env sto t
-          -> Ex (D1E env) (D2 t)  -- the incoming cotangent value
-          -> Ex (Append (D2AcE (Select env sto "accum")) (D1E env)) (TPair (D1 t) (Tup (D2E (Select env sto "merge"))))
-freezeRet descr (Ret e0 subtape e1 sub e2 :: Ret _ _ t) d =
-  let (e0', wInsertD2Ac) = weakenBindings weakenExpr (wSinks (d2ace (select SAccum descr))) e0
+          -> Ex (D2 t : Append (D2AcE (Select env sto "accum")) (D1E env)) (TPair (D1 t) (Tup (D2E (Select env sto "merge"))))
+freezeRet descr (Ret e0 subtape e1 sub e2 :: Ret _ _ t) =
+  let (e0', wInsertD2Ac) = weakenBindings weakenExpr (WSink .> wSinks (d2ace (select SAccum descr))) e0
       e2' = weakenExpr (WCopy (wCopies (subList (bindingsBinds e0) subtape) (wRaiseAbove (d2ace (select SAccum descr)) (sD1eEnv descr)))) e2
   in letBinds e0' $
        EPair ext
          (weakenExpr wInsertD2Ac e1)
-         (ELet ext (weakenExpr (sinkWithBindings e0 .> wSinks (d2ace (select SAccum descr))) d) $
-          ELet ext (weakenExpr (autoWeak (#d (auto1 @(D2 t))
+         (ELet ext (weakenExpr (autoWeak (#d (auto1 @(D2 t))
                                           &. #tape (subList (bindingsBinds e0) subtape)
                                           &. #shbinds (bindingsBinds e0)
                                           &. #d2ace (d2ace (select SAccum descr))
                                           &. #tl (sD1eEnv descr))
                                          (#d :++: LPreW #tape #shbinds (wUndoSubenv subtape) :++: #d2ace :++: #tl)
-                                         (#d :++: #shbinds :++: #d2ace :++: #tl))
+                                         (#shbinds :++: #d :++: #d2ace :++: #tl))
                       e2') $
           expandSubenvZeros (select SMerge descr) sub (EVar ext (tTup (d2e (subList (select SMerge descr) sub))) IZ))
 
diff --git a/src/CHAD/Top.hs b/src/CHAD/Top.hs
new file mode 100644
index 0000000..9df5412
--- /dev/null
+++ b/src/CHAD/Top.hs
@@ -0,0 +1,53 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+module CHAD.Top where
+
+import AST
+import CHAD
+import CHAD.Types
+import Data
+
+
+type family MergeEnv env where
+  MergeEnv '[] = '[]
+  MergeEnv (t : ts) = "merge" : MergeEnv ts
+
+mergeDescr :: SList STy env -> Descr env (MergeEnv env)
+mergeDescr SNil = DTop
+mergeDescr (t `SCons` env) = mergeDescr env `DPush` (t, SMerge)
+
+mergeEnvNoAccum :: SList f env -> Select env (MergeEnv env) "accum" :~: '[]
+mergeEnvNoAccum SNil = Refl
+mergeEnvNoAccum (_ `SCons` env) | Refl <- mergeEnvNoAccum env = Refl
+
+mergeEnvOnlyMerge :: SList f env -> Select env (MergeEnv env) "merge" :~: env
+mergeEnvOnlyMerge SNil = Refl
+mergeEnvOnlyMerge (_ `SCons` env) | Refl <- mergeEnvOnlyMerge env = Refl
+
+d1Identity :: STy t -> D1 t :~: t
+d1Identity = \case
+  STNil -> Refl
+  STPair a b | Refl <- d1Identity a, Refl <- d1Identity b -> Refl
+  STEither a b | Refl <- d1Identity a, Refl <- d1Identity b -> Refl
+  STMaybe t | Refl <- d1Identity t -> Refl
+  STArr _ t | Refl <- d1Identity t -> Refl
+  STScal _ -> Refl
+  STAccum{} -> error "Accumulators not allowed in input program"
+
+d1eIdentity :: SList STy env -> D1E env :~: env
+d1eIdentity SNil = Refl
+d1eIdentity (t `SCons` env) | Refl <- d1Identity t, Refl <- d1eIdentity env = Refl
+
+chad :: SList STy env -> Ex env t -> Ex (D2 t : D1E env) (TPair (D1 t) (Tup (D2E env)))
+chad env term
+  | Refl <- mergeEnvNoAccum env
+  , Refl <- mergeEnvOnlyMerge env
+  = freezeRet (mergeDescr env) (drev (mergeDescr env) term)
+
+chad' :: SList STy env -> Ex env t -> Ex (D2 t : env) (TPair t (Tup (D2E env)))
+chad' env term
+  | Refl <- d1eIdentity env, Refl <- d1Identity (typeOf term)
+  = chad env term
diff --git a/src/Example.hs b/src/Example.hs
index d0405af..1775bb9 100644
--- a/src/Example.hs
+++ b/src/Example.hs
@@ -11,6 +11,7 @@ import Array
 import AST
 import AST.Pretty
 import CHAD
+import CHAD.Top
 import Data
 import ForwardAD
 import Interpreter
@@ -23,16 +24,6 @@ import Example.Format
 
 -- ppExpr senv5 $ simplifyN 20 $ let d = descr5 SMerge SMerge in freezeRet d (drev d ex5) (EConst ext STF32 1.0)
 
-type family MergeEnv env where
-  MergeEnv '[] = '[]
-  MergeEnv (t : ts) = "merge" : MergeEnv ts
-
-mergeDescr :: KnownEnv env => Descr env (MergeEnv env)
-mergeDescr = go knownEnv
-  where go :: SList STy env -> Descr env (MergeEnv env)
-        go SNil = DTop
-        go (t `SCons` env) = go env `DPush` (t, SMerge)
-
 
 bin :: SOp (TPair a b) c -> Ex env a -> Ex env b -> Ex env c
 bin op a b = EOp ext op (EPair ext a b)
@@ -195,9 +186,8 @@ neuralGo =
       argument = (Value input `SCons` Value lay3 `SCons` Value lay2 `SCons` Value lay1 `SCons` SNil)
       revderiv =
         simplifyN 20 $
-        freezeRet mergeDescr
-          (drev mergeDescr neural)
-          (EConst ext STF64 1.0)
+          ELet ext (EConst ext STF64 1.0) $
+            chad knownEnv neural
       (primal, (((((), Right dlay1_1), Right dlay2_1), dlay3_1), dinput_1)) = interpretOpen False argument revderiv
       (Value dinput_2 `SCons` Value dlay3_2 `SCons` Value dlay2_2 `SCons` Value dlay1_2 `SCons` SNil) = drevByFwd knownEnv neural argument 1.0
   in trace (formatter (ppExpr knownEnv revderiv)) $