Initial

1 files changed, 201 insertions, 0 deletions

diff --git a/src/Numeric/ADDual.hs b/src/Numeric/ADDual.hs
new file mode 100644
index 0000000..d9c5b74
--- /dev/null
+++ b/src/Numeric/ADDual.hs
@@ -0,0 +1,201 @@
+{-# LANGUAGE DeriveTraversable #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE MultiWayIf #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+module Numeric.ADDual where
+
+import Control.Monad (when)
+import Control.Monad.Trans.Class (lift)
+import Control.Monad.Trans.State.Strict
+import Data.IORef
+import Data.Proxy
+import qualified Data.Vector.Storable as VS
+import qualified Data.Vector.Storable.Mutable as VSM
+import Foreign.Ptr
+import Foreign.Storable
+import GHC.Exts (withDict)
+import System.IO.Unsafe
+
+-- import System.IO (hPutStrLn, stderr)
+
+-- import Numeric.ADDual.IDGen
+
+
+-- TODO: full vjp (just some more Traversable mess)
+{-# NOINLINE gradient' #-}
+gradient' :: forall a f. (Traversable f, Num a, Storable a)
+          -- => Show a  -- TODO: remove
+          => (forall s. Taping s a => f (Dual s a) -> Dual s a)
+          -> f a -> a -> (a, f a)
+gradient' f inp topctg = unsafePerformIO $ do
+  -- hPutStrLn stderr "Preparing input"
+  let (inp', starti) = runState (traverse (\x -> state (\i -> (Dual x i, i + 1))) inp) 0
+  idref <- newIORef starti
+  vec1 <- VSM.new (max 128 (2 * starti))
+  taperef <- newIORef (MLog idref (Chunk 0 vec1) SLNil)
+
+  -- hPutStrLn stderr "Running function"
+  let Dual result outi = withDict @(Taping () a) taperef $ f @() inp'
+  -- hPutStrLn stderr $ "result = " ++ show result ++ "; outi = " ++ show outi
+  MLog _ lastChunk tapeTail <- readIORef taperef
+
+  -- do tapestr <- showTape (tapeTail `Snoc` lastChunk)
+  --    hPutStrLn stderr $ "tape = " ++ tapestr ""
+
+  -- hPutStrLn stderr "Backpropagating"
+  accums <- VSM.new (outi+1)
+  VSM.write accums outi topctg
+
+  let backpropagate i chunk@(Chunk ci0 vec) tape
+        | i >= ci0 = do
+            ctg <- VSM.read accums i
+            Contrib i1 dx i2 dy <- VSM.read vec (i - ci0)
+            when (i1 /= -1) $ VSM.modify accums (+ ctg*dx) i1
+            when (i2 /= -1) $ VSM.modify accums (+ ctg*dy) i2
+            backpropagate (i-1) chunk tape
+        | otherwise = case tape of
+                        SLNil -> return ()  -- reached end of tape we should loop over
+                        tape'@Snoc{} `Snoc` chunk' -> backpropagate i chunk' tape'
+                        -- When we reach the last chunk, modify it so that its
+                        -- starting index is after the inputs.
+                        SLNil `Snoc` Chunk _ vec' ->
+                          backpropagate i (Chunk starti (VSM.slice starti (VSM.length vec' - starti) vec')) SLNil
+
+  -- Ensure that if there are no more chunks in the tape tail, the starting
+  -- index of the first chunk is adjusted so that backpropagate stops in time.
+  case tapeTail of
+    SLNil -> backpropagate outi (let Chunk _ vec = lastChunk
+                                 in Chunk starti (VSM.slice starti (VSM.length vec - starti) vec))
+                                SLNil
+    Snoc{} -> backpropagate outi lastChunk tapeTail
+
+  -- do accums' <- VS.freeze accums
+  --    hPutStrLn stderr $ "accums = " ++ show accums'
+
+  -- hPutStrLn stderr "Reconstructing gradient"
+  let readDeriv = do i <- get
+                     d <- lift $ VSM.read accums i
+                     put (i+1)
+                     return d
+  grad <- evalStateT (traverse (\_ -> readDeriv) inp) 0
+
+  return (result, grad)
+
+data Snoclist a = SLNil | Snoc !(Snoclist a) !a
+  deriving (Show, Eq, Ord, Functor, Foldable, Traversable)
+
+data Contrib a = Contrib {-# UNPACK #-} !Int a  -- ^ ID == -1 -> no contribution
+                         {-# UNPACK #-} !Int a  -- ^ idem
+  deriving (Show)
+
+instance Storable a => Storable (Contrib a) where
+  sizeOf _ = 2 * (sizeOf (undefined :: Int) + sizeOf (undefined :: a))
+  alignment _ = alignment (undefined :: Int)
+  peek ptr = Contrib <$> peek (castPtr ptr)
+                     <*> peekByteOff (castPtr ptr) (sizeOf (undefined :: Int))
+                     <*> peekByteOff (castPtr ptr) (sizeOf (undefined :: Int) + sizeOf (undefined :: a))
+                     <*> peekByteOff (castPtr ptr) (2 * sizeOf (undefined :: Int) + sizeOf (undefined :: a))
+  poke ptr (Contrib i1 dx i2 dy) = do
+    poke (castPtr ptr) i1
+    pokeByteOff (castPtr ptr) (sizeOf (undefined :: Int)) dx
+    pokeByteOff (castPtr ptr) (sizeOf (undefined :: Int) + sizeOf (undefined :: a)) i2
+    pokeByteOff (castPtr ptr) (2 * sizeOf (undefined :: Int) + sizeOf (undefined :: a)) dy
+
+data Chunk a = Chunk {-# UNPACK #-} !Int  -- ^ First ID in this chunk
+                     {-# UNPACK #-} !(VSM.IOVector (Contrib a))
+
+data MLog s a = MLog !(IORef Int)  -- ^ next ID to generate
+                     {-# UNPACK #-} !(Chunk a)  -- ^ current running chunk
+                     !(Snoclist (Chunk a))  -- ^ tape
+
+showChunk :: (Storable a, Show a) => Chunk a -> IO ShowS
+showChunk (Chunk ci0 vec) = do
+  vec' <- VS.freeze vec
+  return (showString ("Chunk " ++ show ci0 ++ " ") . shows vec')
+
+showTape :: (Storable a, Show a) => Snoclist (Chunk a) -> IO ShowS
+showTape SLNil = return (showString "SLNil")
+showTape (tape `Snoc` chunk) = do
+  s1 <- showTape tape
+  s2 <- showChunk chunk
+  return (s1 . showString " `Snoc` " . s2)
+
+-- | This class does not have any instances defined, on purpose. You'll get one
+-- magically when you differentiate.
+class Taping s a where
+  getTape :: IORef (MLog s a)
+
+data Dual s a = Dual !a
+                     {-# UNPACK #-} !Int  -- ^ -1 if this is a constant
+
+instance Eq a => Eq (Dual s a) where
+  Dual x _ == Dual y _ = x == y
+
+instance Ord a => Ord (Dual s a) where
+  compare (Dual x _) (Dual y _) = compare x y
+
+instance (Num a, Storable a, Taping s a) => Num (Dual s a) where
+  Dual x i1 + Dual y i2 = Dual (x + y)    (writeTape @a (Proxy @s) i1 1 i2 1)
+  Dual x i1 - Dual y i2 = Dual (x - y)    (writeTape @a (Proxy @s) i1 1 i2 (-1))
+  Dual x i1 * Dual y i2 = Dual (x * y)    (writeTape    (Proxy @s) i1 y i2 x)
+  negate (Dual x i1)    = Dual (negate x) (writeTape @a (Proxy @s) i1 (-1) (-1) 0)
+  abs (Dual x i1)       = Dual (abs x)    (writeTape    (Proxy @s) i1 (x * signum x) (-1) 0)
+  signum (Dual x _) = Dual (signum x) (-1)
+  fromInteger n = Dual (fromInteger n) (-1)
+
+data WriteTapeAction a = WTANewvec (VSM.IOVector (Contrib a))
+                       | WTAOldTape (Snoclist (Chunk a))
+
+writeTape :: forall a s proxy. (Num a, Storable a, Taping s a) => proxy s -> Int -> a -> Int -> a -> Int
+writeTape _ i1 dx i2 dy = unsafePerformIO $ writeTapeIO (Proxy @s) i1 dx i2 dy
+
+writeTapeIO :: forall a s proxy. (Num a, Storable a, Taping s a) => proxy s -> Int -> a -> Int -> a -> IO Int
+writeTapeIO _ i1 dx i2 dy = do
+  MLog idref (Chunk ci0 vec) _ <- readIORef (getTape @s)
+  let n = VSM.length vec
+  i <- atomicModifyIORef' idref (\i -> (i + 1, i))
+  let idx = i - ci0
+
+  if | idx < n -> do
+         VSM.write vec idx (Contrib i1 dx i2 dy)
+         return i
+
+     -- check if we'd fit in the next chunk (overwhelmingly likely)
+     | let newlen = 3 * n `div` 2
+     , idx < n + newlen -> do
+         newvec <- VSM.new newlen
+         action <- atomicModifyIORef' (getTape @s) $ \(MLog idref' chunk@(Chunk ci0' vec') tape) ->
+           if | ci0 == ci0' ->
+                 -- Likely (certain when single-threaded): no race condition,
+                 -- we get the chance to put the new chunk in place.
+                 (MLog idref' (Chunk (ci0 + n) newvec) (tape `Snoc` chunk), WTANewvec newvec)
+              | i < ci0' + VSM.length vec' ->
+                 -- Race condition; need to write to appropriate position in this vector.
+                 (MLog idref' chunk tape, WTANewvec vec')
+              | i < ci0' ->
+                 -- Very unlikely; need to write to old chunk in tape.
+                 (MLog idref' chunk tape, WTAOldTape tape)
+              | otherwise ->
+                 -- We got an ID so far in the future that it doesn't even fit
+                 -- in the next chunk. But that can't happen, because we're
+                 -- only in this branch if the ID would have fit in the next
+                 -- chunk in the first place!
+                 error "writeTape: impossible"
+         case action of
+           WTANewvec vec' -> VSM.write vec' (idx - n) (Contrib i1 dx i2 dy)
+           WTAOldTape tape ->
+             let go SLNil = error "writeTape: no appropriate tape chunk?"
+                 go (tape' `Snoc` Chunk ci0' vec')
+                   -- The first comparison here is technically unnecessary, but
+                   -- I'm not courageous enough to remove it.
+                   | ci0' <= i, i < ci0' + VSM.length vec' =
+                       VSM.write vec' (i - ci0') (Contrib i1 dx i2 dy)
+                   | otherwise =
+                       go tape'
+             in go tape
+         return i
+
+     -- there's a tremendous amount of competition, let's just try again
+     | otherwise -> writeTapeIO (Proxy @s) i1 dx i2 dy