{-# LANGUAGE TypeFamilies #-}

module Futhark.AD.Fwd (fwdJVP) where

import Control.Monad
import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.Bifunctor (bimap, second)
import Data.Foldable
import Data.Functor.Product
import Data.List.NonEmpty (NonEmpty (..))
import Data.Map qualified as M
import Data.Tuple (Solo (..), getSolo)
import Futhark.AD.Derivatives
import Futhark.AD.Shared
import Futhark.Analysis.PrimExp.Convert
import Futhark.Builder
import Futhark.IR.SOACS
import Futhark.Tools
import Futhark.Util (interleave, splitAt3, unterleave)

zeroExp :: Type -> Exp SOACS
zeroExp :: Type -> Exp SOACS
zeroExp (Prim PrimType
pt) =
  BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimValue
blankPrimValue PrimType
pt
zeroExp (Array PrimType
pt Shape
shape NoMode
_) =
  BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ Shape -> SubExp -> BasicOp
Replicate Shape
shape (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimValue
blankPrimValue PrimType
pt
zeroExp Type
t = [Char] -> Exp SOACS
forall a. HasCallStack => [Char] -> a
error ([Char] -> Exp SOACS) -> [Char] -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ [Char]
"zeroExp: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall a. Show a => a -> [Char]
show Type
t

class (ArrayShape s) => FromShape s where
  fromShape :: Shape -> s

instance FromShape Shape where
  fromShape :: Shape -> Shape
fromShape = Shape -> Shape
forall a. a -> a
id

instance FromShape ExtShape where
  fromShape :: Shape -> ExtShape
fromShape = (SubExp -> Ext SubExp) -> Shape -> ExtShape
forall a b. (a -> b) -> ShapeBase a -> ShapeBase b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> Ext SubExp
forall a. a -> Ext a
Free

tanType :: (FromShape s, Monoid o) => TypeBase s o -> ADM (TypeBase s o)
tanType :: forall s o.
(FromShape s, Monoid o) =>
TypeBase s o -> ADM (TypeBase s o)
tanType (Acc VName
acc Shape
ispace [Type]
ts) = do
  acc_tan <- VName -> ADM VName
forall a. Tangent a => a -> ADM a
tangent VName
acc
  tan_shape <- askShape
  pure $ Acc acc_tan (tan_shape <> ispace) ts
tanType TypeBase s o
t = do
  shape <- ADM Shape
askShape
  pure $ arrayOf (Prim (elemType t)) (fromShape shape <> arrayShape t) o
  where
    o :: o
o = case TypeBase s o
t of
      Array PrimType
_ s
_ o
o' -> o
o'
      TypeBase s o
_ -> o
forall a. Monoid a => a
mempty

slocal' :: ADM a -> ADM a
slocal' :: forall a. ADM a -> ADM a
slocal' = (RState -> RState) -> ADM a -> ADM a
forall a. (RState -> RState) -> ADM a -> ADM a
slocal RState -> RState
forall a. a -> a
id

slocal :: (RState -> RState) -> ADM a -> ADM a
slocal :: forall a. (RState -> RState) -> ADM a -> ADM a
slocal RState -> RState
f ADM a
m = do
  s <- ADM RState
forall s (m :: * -> *). MonadState s m => m s
get
  modify f
  a <- m
  modify $ \RState
s' -> RState
s' {stateTans = stateTans s}
  pure a

data RState = RState
  { RState -> Map VName VName
stateTans :: M.Map VName VName,
    RState -> VNameSource
stateNameSource :: VNameSource
  }

data FEnv = FEnv
  { FEnv -> Shape
envTanShape :: Shape,
    FEnv -> Attrs
envAttrs :: Attrs
  }

newtype ADM a = ADM (BuilderT SOACS (ReaderT FEnv (State RState)) a)
  deriving
    ( (forall a b. (a -> b) -> ADM a -> ADM b)
-> (forall a b. a -> ADM b -> ADM a) -> Functor ADM
forall a b. a -> ADM b -> ADM a
forall a b. (a -> b) -> ADM a -> ADM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> ADM a -> ADM b
fmap :: forall a b. (a -> b) -> ADM a -> ADM b
$c<$ :: forall a b. a -> ADM b -> ADM a
<$ :: forall a b. a -> ADM b -> ADM a
Functor,
      Functor ADM
Functor ADM =>
(forall a. a -> ADM a)
-> (forall a b. ADM (a -> b) -> ADM a -> ADM b)
-> (forall a b c. (a -> b -> c) -> ADM a -> ADM b -> ADM c)
-> (forall a b. ADM a -> ADM b -> ADM b)
-> (forall a b. ADM a -> ADM b -> ADM a)
-> Applicative ADM
forall a. a -> ADM a
forall a b. ADM a -> ADM b -> ADM a
forall a b. ADM a -> ADM b -> ADM b
forall a b. ADM (a -> b) -> ADM a -> ADM b
forall a b c. (a -> b -> c) -> ADM a -> ADM b -> ADM c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> ADM a
pure :: forall a. a -> ADM a
$c<*> :: forall a b. ADM (a -> b) -> ADM a -> ADM b
<*> :: forall a b. ADM (a -> b) -> ADM a -> ADM b
$cliftA2 :: forall a b c. (a -> b -> c) -> ADM a -> ADM b -> ADM c
liftA2 :: forall a b c. (a -> b -> c) -> ADM a -> ADM b -> ADM c
$c*> :: forall a b. ADM a -> ADM b -> ADM b
*> :: forall a b. ADM a -> ADM b -> ADM b
$c<* :: forall a b. ADM a -> ADM b -> ADM a
<* :: forall a b. ADM a -> ADM b -> ADM a
Applicative,
      Applicative ADM
Applicative ADM =>
(forall a b. ADM a -> (a -> ADM b) -> ADM b)
-> (forall a b. ADM a -> ADM b -> ADM b)
-> (forall a. a -> ADM a)
-> Monad ADM
forall a. a -> ADM a
forall a b. ADM a -> ADM b -> ADM b
forall a b. ADM a -> (a -> ADM b) -> ADM b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b. ADM a -> (a -> ADM b) -> ADM b
>>= :: forall a b. ADM a -> (a -> ADM b) -> ADM b
$c>> :: forall a b. ADM a -> ADM b -> ADM b
>> :: forall a b. ADM a -> ADM b -> ADM b
$creturn :: forall a. a -> ADM a
return :: forall a. a -> ADM a
Monad,
      MonadState RState,
      MonadReader FEnv,
      Monad ADM
ADM VNameSource
Monad ADM =>
ADM VNameSource -> (VNameSource -> ADM ()) -> MonadFreshNames ADM
VNameSource -> ADM ()
forall (m :: * -> *).
Monad m =>
m VNameSource -> (VNameSource -> m ()) -> MonadFreshNames m
$cgetNameSource :: ADM VNameSource
getNameSource :: ADM VNameSource
$cputNameSource :: VNameSource -> ADM ()
putNameSource :: VNameSource -> ADM ()
MonadFreshNames,
      HasScope SOACS,
      LocalScope SOACS
    )

instance MonadBuilder ADM where
  type Rep ADM = SOACS
  mkExpDecM :: Pat (LetDec (Rep ADM)) -> Exp (Rep ADM) -> ADM (ExpDec (Rep ADM))
mkExpDecM Pat (LetDec (Rep ADM))
pat Exp (Rep ADM)
e = BuilderT SOACS (ReaderT FEnv (State RState)) (ExpDec (Rep ADM))
-> ADM (ExpDec (Rep ADM))
forall a. BuilderT SOACS (ReaderT FEnv (State RState)) a -> ADM a
ADM (BuilderT SOACS (ReaderT FEnv (State RState)) (ExpDec (Rep ADM))
 -> ADM (ExpDec (Rep ADM)))
-> BuilderT SOACS (ReaderT FEnv (State RState)) (ExpDec (Rep ADM))
-> ADM (ExpDec (Rep ADM))
forall a b. (a -> b) -> a -> b
$ Pat (LetDec (Rep (BuilderT SOACS (ReaderT FEnv (State RState)))))
-> Exp (Rep (BuilderT SOACS (ReaderT FEnv (State RState))))
-> BuilderT
     SOACS
     (ReaderT FEnv (State RState))
     (ExpDec (Rep (BuilderT SOACS (ReaderT FEnv (State RState)))))
forall (m :: * -> *).
MonadBuilder m =>
Pat (LetDec (Rep m)) -> Exp (Rep m) -> m (ExpDec (Rep m))
mkExpDecM Pat (LetDec (Rep (BuilderT SOACS (ReaderT FEnv (State RState)))))
Pat (LetDec (Rep ADM))
pat Exp (Rep (BuilderT SOACS (ReaderT FEnv (State RState))))
Exp (Rep ADM)
e
  mkBodyM :: forall res.
IsResult res =>
Stms (Rep ADM) -> [res] -> ADM (GBody (Rep ADM) res)
mkBodyM Stms (Rep ADM)
bnds [res]
res = BuilderT SOACS (ReaderT FEnv (State RState)) (GBody (Rep ADM) res)
-> ADM (GBody (Rep ADM) res)
forall a. BuilderT SOACS (ReaderT FEnv (State RState)) a -> ADM a
ADM (BuilderT SOACS (ReaderT FEnv (State RState)) (GBody (Rep ADM) res)
 -> ADM (GBody (Rep ADM) res))
-> BuilderT
     SOACS (ReaderT FEnv (State RState)) (GBody (Rep ADM) res)
-> ADM (GBody (Rep ADM) res)
forall a b. (a -> b) -> a -> b
$ Stms (Rep (BuilderT SOACS (ReaderT FEnv (State RState))))
-> [res]
-> BuilderT
     SOACS
     (ReaderT FEnv (State RState))
     (GBody (Rep (BuilderT SOACS (ReaderT FEnv (State RState)))) res)
forall res.
IsResult res =>
Stms (Rep (BuilderT SOACS (ReaderT FEnv (State RState))))
-> [res]
-> BuilderT
     SOACS
     (ReaderT FEnv (State RState))
     (GBody (Rep (BuilderT SOACS (ReaderT FEnv (State RState)))) res)
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
Stms (Rep m) -> [res] -> m (GBody (Rep m) res)
mkBodyM Stms (Rep (BuilderT SOACS (ReaderT FEnv (State RState))))
Stms (Rep ADM)
bnds [res]
res
  mkLetNamesM :: [VName] -> Exp (Rep ADM) -> ADM (Stm (Rep ADM))
mkLetNamesM [VName]
pat Exp (Rep ADM)
e = BuilderT SOACS (ReaderT FEnv (State RState)) (Stm (Rep ADM))
-> ADM (Stm (Rep ADM))
forall a. BuilderT SOACS (ReaderT FEnv (State RState)) a -> ADM a
ADM (BuilderT SOACS (ReaderT FEnv (State RState)) (Stm (Rep ADM))
 -> ADM (Stm (Rep ADM)))
-> BuilderT SOACS (ReaderT FEnv (State RState)) (Stm (Rep ADM))
-> ADM (Stm (Rep ADM))
forall a b. (a -> b) -> a -> b
$ [VName]
-> Exp (Rep (BuilderT SOACS (ReaderT FEnv (State RState))))
-> BuilderT
     SOACS
     (ReaderT FEnv (State RState))
     (Stm (Rep (BuilderT SOACS (ReaderT FEnv (State RState)))))
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m (Stm (Rep m))
mkLetNamesM [VName]
pat Exp (Rep (BuilderT SOACS (ReaderT FEnv (State RState))))
Exp (Rep ADM)
e

  addStms :: Stms (Rep ADM) -> ADM ()
addStms = BuilderT SOACS (ReaderT FEnv (State RState)) () -> ADM ()
forall a. BuilderT SOACS (ReaderT FEnv (State RState)) a -> ADM a
ADM (BuilderT SOACS (ReaderT FEnv (State RState)) () -> ADM ())
-> (Stms SOACS -> BuilderT SOACS (ReaderT FEnv (State RState)) ())
-> Stms SOACS
-> ADM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms (Rep (BuilderT SOACS (ReaderT FEnv (State RState))))
-> BuilderT SOACS (ReaderT FEnv (State RState)) ()
Stms SOACS -> BuilderT SOACS (ReaderT FEnv (State RState)) ()
forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms
  collectStms :: forall a. ADM a -> ADM (a, Stms (Rep ADM))
collectStms (ADM BuilderT SOACS (ReaderT FEnv (State RState)) a
m) = BuilderT SOACS (ReaderT FEnv (State RState)) (a, Stms (Rep ADM))
-> ADM (a, Stms (Rep ADM))
forall a. BuilderT SOACS (ReaderT FEnv (State RState)) a -> ADM a
ADM (BuilderT SOACS (ReaderT FEnv (State RState)) (a, Stms (Rep ADM))
 -> ADM (a, Stms (Rep ADM)))
-> BuilderT SOACS (ReaderT FEnv (State RState)) (a, Stms (Rep ADM))
-> ADM (a, Stms (Rep ADM))
forall a b. (a -> b) -> a -> b
$ BuilderT SOACS (ReaderT FEnv (State RState)) a
-> BuilderT
     SOACS
     (ReaderT FEnv (State RState))
     (a, Stms (Rep (BuilderT SOACS (ReaderT FEnv (State RState)))))
forall a.
BuilderT SOACS (ReaderT FEnv (State RState)) a
-> BuilderT
     SOACS
     (ReaderT FEnv (State RState))
     (a, Stms (Rep (BuilderT SOACS (ReaderT FEnv (State RState)))))
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms BuilderT SOACS (ReaderT FEnv (State RState)) a
m

instance MonadFreshNames (State RState) where
  getNameSource :: State RState VNameSource
getNameSource = (RState -> VNameSource) -> State RState VNameSource
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets RState -> VNameSource
stateNameSource
  putNameSource :: VNameSource -> State RState ()
putNameSource VNameSource
src = (RState -> RState) -> State RState ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (\RState
env -> RState
env {stateNameSource = src})

askShape :: ADM Shape
askShape :: ADM Shape
askShape = BuilderT SOACS (ReaderT FEnv (State RState)) Shape -> ADM Shape
forall a. BuilderT SOACS (ReaderT FEnv (State RState)) a -> ADM a
ADM (BuilderT SOACS (ReaderT FEnv (State RState)) Shape -> ADM Shape)
-> BuilderT SOACS (ReaderT FEnv (State RState)) Shape -> ADM Shape
forall a b. (a -> b) -> a -> b
$ ReaderT FEnv (State RState) Shape
-> BuilderT SOACS (ReaderT FEnv (State RState)) Shape
forall (m :: * -> *) a. Monad m => m a -> BuilderT SOACS m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ReaderT FEnv (State RState) Shape
 -> BuilderT SOACS (ReaderT FEnv (State RState)) Shape)
-> ReaderT FEnv (State RState) Shape
-> BuilderT SOACS (ReaderT FEnv (State RState)) Shape
forall a b. (a -> b) -> a -> b
$ (FEnv -> Shape) -> ReaderT FEnv (State RState) Shape
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks FEnv -> Shape
envTanShape

runADM :: (MonadFreshNames m) => Shape -> Attrs -> ADM a -> m a
runADM :: forall (m :: * -> *) a.
MonadFreshNames m =>
Shape -> Attrs -> ADM a -> m a
runADM Shape
shape Attrs
attrs (ADM BuilderT SOACS (ReaderT FEnv (State RState)) a
m) =
  (VNameSource -> (a, VNameSource)) -> m a
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> (a, VNameSource)) -> m a)
-> (VNameSource -> (a, VNameSource)) -> m a
forall a b. (a -> b) -> a -> b
$ \VNameSource
vn ->
    (RState -> VNameSource) -> (a, RState) -> (a, VNameSource)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second RState -> VNameSource
stateNameSource ((a, RState) -> (a, VNameSource))
-> (a, RState) -> (a, VNameSource)
forall a b. (a -> b) -> a -> b
$
      State RState a -> RState -> (a, RState)
forall s a. State s a -> s -> (a, s)
runState
        ( ReaderT FEnv (State RState) a -> FEnv -> State RState a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT
            ((a, Stms SOACS) -> a
forall a b. (a, b) -> a
fst ((a, Stms SOACS) -> a)
-> ReaderT FEnv (State RState) (a, Stms SOACS)
-> ReaderT FEnv (State RState) a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BuilderT SOACS (ReaderT FEnv (State RState)) a
-> Scope SOACS -> ReaderT FEnv (State RState) (a, Stms SOACS)
forall (m :: * -> *) rep a.
MonadFreshNames m =>
BuilderT rep m a -> Scope rep -> m (a, Stms rep)
runBuilderT BuilderT SOACS (ReaderT FEnv (State RState)) a
m Scope SOACS
forall a. Monoid a => a
mempty)
            (Shape -> Attrs -> FEnv
FEnv Shape
shape Attrs
attrs)
        )
        (Map VName VName -> VNameSource -> RState
RState Map VName VName
forall a. Monoid a => a
mempty VNameSource
vn)

tanVName :: VName -> ADM VName
tanVName :: VName -> ADM VName
tanVName VName
v = Name -> ADM VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName (VName -> Name
baseName VName
v Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_tan")

insertTan :: VName -> VName -> ADM ()
insertTan :: VName -> VName -> ADM ()
insertTan VName
v VName
v' =
  (RState -> RState) -> ADM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((RState -> RState) -> ADM ()) -> (RState -> RState) -> ADM ()
forall a b. (a -> b) -> a -> b
$ \RState
env -> RState
env {stateTans = M.insert v v' (stateTans env)}

class TanBuilder a where
  newTan :: a -> ADM a
  bundleNew :: a -> ADM (a, a)

bundleNewList :: (TanBuilder a) => [a] -> ADM [a]
bundleNewList :: forall a. TanBuilder a => [a] -> ADM [a]
bundleNewList = ([(a, a)] -> [a]) -> ADM [(a, a)] -> ADM [a]
forall a b. (a -> b) -> ADM a -> ADM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([a] -> [a] -> [a]) -> ([a], [a]) -> [a]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
interleave (([a], [a]) -> [a]) -> ([(a, a)] -> ([a], [a])) -> [(a, a)] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(a, a)] -> ([a], [a])
forall a b. [(a, b)] -> ([a], [b])
unzip) (ADM [(a, a)] -> ADM [a])
-> ([a] -> ADM [(a, a)]) -> [a] -> ADM [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> ADM (a, a)) -> [a] -> ADM [(a, a)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM a -> ADM (a, a)
forall a. TanBuilder a => a -> ADM (a, a)
bundleNew

instance (FromShape s, Monoid o) => TanBuilder (PatElem (TypeBase s o)) where
  newTan :: PatElem (TypeBase s o) -> ADM (PatElem (TypeBase s o))
newTan (PatElem VName
p TypeBase s o
t) = do
    p' <- VName -> ADM VName
tanVName VName
p
    insertTan p p'
    t' <- tanType t
    pure $ PatElem p' t'
  bundleNew :: PatElem (TypeBase s o)
-> ADM (PatElem (TypeBase s o), PatElem (TypeBase s o))
bundleNew PatElem (TypeBase s o)
pe = do
    pe' <- PatElem (TypeBase s o) -> ADM (PatElem (TypeBase s o))
forall a. TanBuilder a => a -> ADM a
newTan PatElem (TypeBase s o)
pe
    pure (pe, pe')

newTanPat :: (TanBuilder (PatElem t)) => Pat t -> ADM (Pat t)
newTanPat :: forall t. TanBuilder (PatElem t) => Pat t -> ADM (Pat t)
newTanPat (Pat [PatElem t]
pes) = [PatElem t] -> Pat t
forall dec. [PatElem dec] -> Pat dec
Pat ([PatElem t] -> Pat t) -> ADM [PatElem t] -> ADM (Pat t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatElem t -> ADM (PatElem t)) -> [PatElem t] -> ADM [PatElem t]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM PatElem t -> ADM (PatElem t)
forall a. TanBuilder a => a -> ADM a
newTan [PatElem t]
pes

bundleNewPat :: (TanBuilder (PatElem t)) => Pat t -> ADM (Pat t)
bundleNewPat :: forall t. TanBuilder (PatElem t) => Pat t -> ADM (Pat t)
bundleNewPat (Pat [PatElem t]
pes) = [PatElem t] -> Pat t
forall dec. [PatElem dec] -> Pat dec
Pat ([PatElem t] -> Pat t) -> ADM [PatElem t] -> ADM (Pat t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [PatElem t] -> ADM [PatElem t]
forall a. TanBuilder a => [a] -> ADM [a]
bundleNewList [PatElem t]
pes

instance (FromShape s, Monoid o) => TanBuilder (Param (TypeBase s o)) where
  newTan :: Param (TypeBase s o) -> ADM (Param (TypeBase s o))
newTan (Param Attrs
_ VName
p TypeBase s o
t) = do
    PatElem p' t' <- PatElem (TypeBase s o) -> ADM (PatElem (TypeBase s o))
forall a. TanBuilder a => a -> ADM a
newTan (PatElem (TypeBase s o) -> ADM (PatElem (TypeBase s o)))
-> PatElem (TypeBase s o) -> ADM (PatElem (TypeBase s o))
forall a b. (a -> b) -> a -> b
$ VName -> TypeBase s o -> PatElem (TypeBase s o)
forall dec. VName -> dec -> PatElem dec
PatElem VName
p TypeBase s o
t
    pure $ Param mempty p' t'
  bundleNew :: Param (TypeBase s o)
-> ADM (Param (TypeBase s o), Param (TypeBase s o))
bundleNew Param (TypeBase s o)
param = do
    param' <- Param (TypeBase s o) -> ADM (Param (TypeBase s o))
forall a. TanBuilder a => a -> ADM a
newTan Param (TypeBase s o)
param
    pure (param, param')

instance (TanBuilder a, Tangent b) => TanBuilder (a, b) where
  newTan :: (a, b) -> ADM (a, b)
newTan (a
p, b
x) = (,) (a -> b -> (a, b)) -> ADM a -> ADM (b -> (a, b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> ADM a
forall a. TanBuilder a => a -> ADM a
newTan a
p ADM (b -> (a, b)) -> ADM b -> ADM (a, b)
forall a b. ADM (a -> b) -> ADM a -> ADM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> b -> ADM b
forall a. Tangent a => a -> ADM a
tangent b
x
  bundleNew :: (a, b) -> ADM ((a, b), (a, b))
bundleNew (a
p, b
x) = do
    p' <- a -> ADM a
forall a. TanBuilder a => a -> ADM a
newTan a
p
    x_tan <- tangent x
    pure ((p, x), (p', x_tan))

class Tangent a where
  tangent :: a -> ADM a
  bundleTan :: a -> ADM (a, a)

instance (FromShape s, Monoid o) => Tangent (TypeBase s o) where
  tangent :: TypeBase s o -> ADM (TypeBase s o)
tangent = TypeBase s o -> ADM (TypeBase s o)
forall s o.
(FromShape s, Monoid o) =>
TypeBase s o -> ADM (TypeBase s o)
tanType
  bundleTan :: TypeBase s o -> ADM (TypeBase s o, TypeBase s o)
bundleTan TypeBase s o
t = do
    t' <- TypeBase s o -> ADM (TypeBase s o)
forall a. Tangent a => a -> ADM a
tangent TypeBase s o
t
    pure (t, t')

bundleTangents :: (Tangent a) => [a] -> ADM [a]
bundleTangents :: forall a. Tangent a => [a] -> ADM [a]
bundleTangents = ([(a, a)] -> [a]) -> ADM [(a, a)] -> ADM [a]
forall a b. (a -> b) -> ADM a -> ADM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([a] -> [a] -> [a]) -> ([a], [a]) -> [a]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
interleave (([a], [a]) -> [a]) -> ([(a, a)] -> ([a], [a])) -> [(a, a)] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(a, a)] -> ([a], [a])
forall a b. [(a, b)] -> ([a], [b])
unzip) (ADM [(a, a)] -> ADM [a])
-> ([a] -> ADM [(a, a)]) -> [a] -> ADM [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> ADM (a, a)) -> [a] -> ADM [(a, a)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM a -> ADM (a, a)
forall a. Tangent a => a -> ADM (a, a)
bundleTan

instance Tangent VName where
  tangent :: VName -> ADM VName
tangent VName
v = do
    maybeTan <- (RState -> Maybe VName) -> ADM (Maybe VName)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((RState -> Maybe VName) -> ADM (Maybe VName))
-> (RState -> Maybe VName) -> ADM (Maybe VName)
forall a b. (a -> b) -> a -> b
$ VName -> Map VName VName -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (Map VName VName -> Maybe VName)
-> (RState -> Map VName VName) -> RState -> Maybe VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RState -> Map VName VName
stateTans
    case maybeTan of
      Just VName
v_tan -> VName -> ADM VName
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v_tan
      Maybe VName
Nothing -> do
        t <- VName -> ADM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
        when (isAcc t) $
          error $
            "Missing tangent for accumulator " <> prettyString v
        tan_shape <- askShape
        letExp (baseName v <> "_implicit_tan") $ zeroExp $ t `arrayOfShape` tan_shape
  bundleTan :: VName -> ADM (VName, VName)
bundleTan VName
v = do
    v_tan <- VName -> ADM VName
forall a. Tangent a => a -> ADM a
tangent VName
v
    pure (v, v_tan)

instance Tangent SubExp where
  tangent :: SubExp -> ADM SubExp
tangent (Constant PrimValue
c) = do
    tan_shape <- ADM Shape
askShape
    if tan_shape == mempty
      then pure $ constant $ blankPrimValue pt
      else letSubExp "const_implicit_tan" $ zeroExp $ Prim pt `arrayOfShape` tan_shape
    where
      pt :: PrimType
pt = PrimValue -> PrimType
primValueType PrimValue
c
  tangent (Var VName
v) = VName -> SubExp
Var (VName -> SubExp) -> ADM VName -> ADM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> ADM VName
forall a. Tangent a => a -> ADM a
tangent VName
v
  bundleTan :: SubExp -> ADM (SubExp, SubExp)
bundleTan c :: SubExp
c@Constant {} = do
    c_tan <- SubExp -> ADM SubExp
forall a. Tangent a => a -> ADM a
tangent SubExp
c
    pure (c, c_tan)
  bundleTan (Var VName
v) = (VName -> SubExp)
-> (VName -> SubExp) -> (VName, VName) -> (SubExp, SubExp)
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap VName -> SubExp
Var VName -> SubExp
Var ((VName, VName) -> (SubExp, SubExp))
-> ADM (VName, VName) -> ADM (SubExp, SubExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> ADM (VName, VName)
forall a. Tangent a => a -> ADM (a, a)
bundleTan VName
v

instance Tangent SubExpRes where
  tangent :: SubExpRes -> ADM SubExpRes
tangent (SubExpRes Certs
cs SubExp
se) = Certs -> SubExp -> SubExpRes
SubExpRes Certs
cs (SubExp -> SubExpRes) -> ADM SubExp -> ADM SubExpRes
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> ADM SubExp
forall a. Tangent a => a -> ADM a
tangent SubExp
se
  bundleTan :: SubExpRes -> ADM (SubExpRes, SubExpRes)
bundleTan (SubExpRes Certs
cs SubExp
se) = (SubExp -> SubExpRes)
-> (SubExp -> SubExpRes)
-> (SubExp, SubExp)
-> (SubExpRes, SubExpRes)
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (Certs -> SubExp -> SubExpRes
SubExpRes Certs
cs) (Certs -> SubExp -> SubExpRes
SubExpRes Certs
cs) ((SubExp, SubExp) -> (SubExpRes, SubExpRes))
-> ADM (SubExp, SubExp) -> ADM (SubExpRes, SubExpRes)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> ADM (SubExp, SubExp)
forall a. Tangent a => a -> ADM (a, a)
bundleTan SubExp
se

withTan ::
  SubExp ->
  (SubExp -> ADM (Exp SOACS)) ->
  ADM (Exp SOACS)
withTan :: SubExp -> (SubExp -> ADM (Exp SOACS)) -> ADM (Exp SOACS)
withTan SubExp
x SubExp -> ADM (Exp SOACS)
f = do
  shape <- ADM Shape
askShape
  x_tan <- tangent x
  mapNest shape (MkSolo x_tan) (f . getSolo)

withTansI ::
  VName ->
  [SubExp] ->
  ([SubExp] -> VName -> [SubExp] -> ADM (Exp SOACS)) ->
  ADM (Exp SOACS)
withTansI :: VName
-> [SubExp]
-> ([SubExp] -> VName -> [SubExp] -> ADM (Exp SOACS))
-> ADM (Exp SOACS)
withTansI VName
x [SubExp]
ys [SubExp] -> VName -> [SubExp] -> ADM (Exp SOACS)
f = do
  shape <- ADM Shape
askShape
  x_tan <- tangent x
  ys_tan <- mapM tangent ys
  if shape == mempty
    then f [] x_tan ys_tan
    else do
      let w = Int -> Shape -> SubExp
shapeSize Int
0 Shape
shape
      ys_tan_vs <- mapM asVName ys_tan
      iota_p <- newParam "iota_p" $ Prim int64
      x_tan_p <- newParam "x_tanp" . rowType =<< lookupType x_tan
      ys_tan_ps <- mapM (newParam "y_tanp" . rowType <=< lookupType) ys_tan_vs
      lam <- mkLambda (iota_p : x_tan_p : ys_tan_ps) $ do
        fmap (subExpsRes . pure) . letSubExp "tan"
          =<< f
            [Var $ paramName iota_p]
            (paramName x_tan_p)
            (map (Var . paramName) ys_tan_ps)
      iota_v <- letExp "iota" $ iota64 w
      Op . Screma w (iota_v : x_tan : ys_tan_vs) <$> mapSOAC lam

withTans ::
  PrimType ->
  SubExp ->
  SubExp ->
  (PrimExp VName -> PrimExp VName -> PrimExp VName) ->
  ADM (Exp SOACS)
withTans :: PrimType
-> SubExp
-> SubExp
-> (PrimExp VName -> PrimExp VName -> PrimExp VName)
-> ADM (Exp SOACS)
withTans PrimType
t SubExp
x SubExp
y PrimExp VName -> PrimExp VName -> PrimExp VName
f = do
  shape <- ADM Shape
askShape
  x_tan <- tangent x
  y_tan <- tangent y
  mapNest shape (Pair (Identity x_tan) (Identity y_tan)) $ \Product Identity Identity SubExp
xy -> do
    Pair (Identity x_tan_v) (Identity y_tan_v) <- (SubExp -> ADM VName)
-> Product Identity Identity SubExp
-> ADM (Product Identity Identity VName)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b)
-> Product Identity Identity a -> f (Product Identity Identity b)
traverse SubExp -> ADM VName
forall (m :: * -> *). MonadBuilder m => SubExp -> m VName
asVName Product Identity Identity SubExp
xy
    toExp $ f (LeafExp x_tan_v t) (LeafExp y_tan_v t)

withAnyTans ::
  (Traversable f) =>
  f SubExp ->
  ([PrimExp VName] -> PrimExp VName) ->
  ADM (Exp SOACS)
withAnyTans :: forall (f :: * -> *).
Traversable f =>
f SubExp -> ([PrimExp VName] -> PrimExp VName) -> ADM (Exp SOACS)
withAnyTans f SubExp
xs [PrimExp VName] -> PrimExp VName
f = do
  shape <- ADM Shape
askShape
  xs_tan <- traverse tangent xs
  mapNest shape xs_tan $ \f SubExp
xs_tan' -> do
    xs_tan'' <- f SubExp
-> (SubExp -> ADM (PrimExp VName)) -> ADM (f (PrimExp VName))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM f SubExp
xs_tan' ((SubExp -> ADM (PrimExp VName)) -> ADM (f (PrimExp VName)))
-> (SubExp -> ADM (PrimExp VName)) -> ADM (f (PrimExp VName))
forall a b. (a -> b) -> a -> b
$ \SubExp
se -> do
      ~(Prim t) <- SubExp -> ADM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
se
      pure $ primExpFromSubExp t se
    toExp $ f $ toList xs_tan''

bindTanPat :: Pat Type -> StmAux () -> Exp SOACS -> ADM ()
bindTanPat :: Pat Type -> StmAux () -> Exp SOACS -> ADM ()
bindTanPat Pat Type
pat_tan StmAux ()
aux Exp SOACS
e = do
  attrs <- (FEnv -> Attrs) -> ADM Attrs
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks FEnv -> Attrs
envAttrs
  auxing aux . attributing attrs . letBind pat_tan $ e

bindTan ::
  Pat Type ->
  StmAux () ->
  SubExp ->
  (SubExp -> ADM (Exp SOACS)) ->
  ADM ()
bindTan :: Pat Type
-> StmAux () -> SubExp -> (SubExp -> ADM (Exp SOACS)) -> ADM ()
bindTan Pat Type
pat_tan StmAux ()
aux SubExp
x SubExp -> ADM (Exp SOACS)
f = do
  Pat Type -> StmAux () -> Exp SOACS -> ADM ()
bindTanPat Pat Type
pat_tan StmAux ()
aux (Exp SOACS -> ADM ()) -> ADM (Exp SOACS) -> ADM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SubExp -> (SubExp -> ADM (Exp SOACS)) -> ADM (Exp SOACS)
withTan SubExp
x SubExp -> ADM (Exp SOACS)
f

bindTans ::
  Pat Type ->
  StmAux () ->
  PrimType ->
  SubExp ->
  SubExp ->
  (PrimExp VName -> PrimExp VName -> PrimExp VName) ->
  ADM ()
bindTans :: Pat Type
-> StmAux ()
-> PrimType
-> SubExp
-> SubExp
-> (PrimExp VName -> PrimExp VName -> PrimExp VName)
-> ADM ()
bindTans Pat Type
pat_tan StmAux ()
aux PrimType
t SubExp
x SubExp
y PrimExp VName -> PrimExp VName -> PrimExp VName
f = do
  Pat Type -> StmAux () -> Exp SOACS -> ADM ()
bindTanPat Pat Type
pat_tan StmAux ()
aux (Exp SOACS -> ADM ()) -> ADM (Exp SOACS) -> ADM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PrimType
-> SubExp
-> SubExp
-> (PrimExp VName -> PrimExp VName -> PrimExp VName)
-> ADM (Exp SOACS)
withTans PrimType
t SubExp
x SubExp
y PrimExp VName -> PrimExp VName -> PrimExp VName
f

basicFwd :: Pat Type -> StmAux () -> BasicOp -> ADM ()
basicFwd :: Pat Type -> StmAux () -> BasicOp -> ADM ()
basicFwd Pat Type
pat StmAux ()
aux BasicOp
op = do
  pat_tan <- Pat Type -> ADM (Pat Type)
forall t. TanBuilder (PatElem t) => Pat t -> ADM (Pat t)
newTanPat Pat Type
pat
  case op of
    SubExp SubExp
se -> do
      se_tan <- SubExp -> ADM SubExp
forall a. Tangent a => a -> ADM a
tangent SubExp
se
      addStm $ Let pat_tan aux $ BasicOp $ SubExp se_tan
    Opaque OpaqueOp
opaqueop SubExp
se -> do
      se_tan <- SubExp -> ADM SubExp
forall a. Tangent a => a -> ADM a
tangent SubExp
se
      addStm $ Let pat_tan aux $ BasicOp $ Opaque opaqueop se_tan
    ArrayLit [SubExp]
ses Type
t -> do
      tan_shape <- ADM Shape
askShape
      ses_tan <- mapM tangent ses
      if tan_shape == mempty
        then
          addStm $ Let pat_tan aux $ BasicOp $ ArrayLit ses_tan t
        else do
          pat_tan_tr <- letExp "pat_tan_tr" $ BasicOp $ ArrayLit ses_tan $ t `arrayOfShape` tan_shape
          pat_tan_tr_t <- lookupType pat_tan_tr
          let perm = Shape -> Type -> [Int]
forall s o. ArrayShape s => Shape -> TypeBase s o -> [Int]
vecPerm Shape
tan_shape Type
pat_tan_tr_t
          addStm $ Let pat_tan aux $ BasicOp $ Rearrange pat_tan_tr perm
    UnOp UnOp
unop SubExp
x -> do
      let t :: PrimType
t = UnOp -> PrimType
unOpType UnOp
unop
          x_pe :: PrimExp VName
x_pe = PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
t SubExp
x
          dx :: PrimExp VName
dx = UnOp -> PrimExp VName -> PrimExp VName
pdUnOp UnOp
unop PrimExp VName
x_pe
      Pat Type
-> StmAux () -> SubExp -> (SubExp -> ADM (Exp SOACS)) -> ADM ()
bindTan Pat Type
pat_tan StmAux ()
aux SubExp
x ((SubExp -> ADM (Exp SOACS)) -> ADM ())
-> (SubExp -> ADM (Exp SOACS)) -> ADM ()
forall a b. (a -> b) -> a -> b
$ \SubExp
x_tan ->
        PrimExp VName -> ADM (Exp (Rep ADM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
PrimExp VName -> m (Exp (Rep m))
toExp (PrimExp VName -> ADM (Exp (Rep ADM)))
-> PrimExp VName -> ADM (Exp (Rep ADM))
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
t SubExp
x_tan PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. PrimExp v -> PrimExp v -> PrimExp v
~*~ PrimExp VName
dx
    BinOp BinOp
bop SubExp
x SubExp
y -> do
      let t :: PrimType
t = BinOp -> PrimType
binOpType BinOp
bop
      Pat Type
-> StmAux ()
-> PrimType
-> SubExp
-> SubExp
-> (PrimExp VName -> PrimExp VName -> PrimExp VName)
-> ADM ()
bindTans Pat Type
pat_tan StmAux ()
aux PrimType
t SubExp
x SubExp
y ((PrimExp VName -> PrimExp VName -> PrimExp VName) -> ADM ())
-> (PrimExp VName -> PrimExp VName -> PrimExp VName) -> ADM ()
forall a b. (a -> b) -> a -> b
$ \PrimExp VName
x_tan PrimExp VName
y_tan ->
        let (PrimExp VName
wrt_x, PrimExp VName
wrt_y) =
              BinOp
-> PrimExp VName -> PrimExp VName -> (PrimExp VName, PrimExp VName)
pdBinOp BinOp
bop (PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
t SubExp
x) (PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
t SubExp
y)
         in PrimExp VName
x_tan PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. PrimExp v -> PrimExp v -> PrimExp v
~*~ PrimExp VName
wrt_x PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. PrimExp v -> PrimExp v -> PrimExp v
~+~ PrimExp VName
y_tan PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. PrimExp v -> PrimExp v -> PrimExp v
~*~ PrimExp VName
wrt_y
    CmpOp {} -> do
      tan_shape <- ADM Shape
askShape
      addStm $ Let pat_tan aux $ zeroExp $ Prim Bool `arrayOfShape` tan_shape
    ConvOp ConvOp
cop SubExp
x ->
      Pat Type
-> StmAux () -> SubExp -> (SubExp -> ADM (Exp SOACS)) -> ADM ()
bindTan Pat Type
pat_tan StmAux ()
aux SubExp
x ((SubExp -> ADM (Exp SOACS)) -> ADM ())
-> (SubExp -> ADM (Exp SOACS)) -> ADM ()
forall a b. (a -> b) -> a -> b
$ \SubExp
x_tan ->
        Exp SOACS -> ADM (Exp SOACS)
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp SOACS -> ADM (Exp SOACS)) -> Exp SOACS -> ADM (Exp SOACS)
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
ConvOp ConvOp
cop SubExp
x_tan
    Assert {} -> () -> ADM ()
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Index VName
arr Slice SubExp
slice -> do
      dims <- Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims (Shape -> [SubExp]) -> ADM Shape -> ADM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ADM Shape
askShape
      arr_tan <- tangent arr
      let slice' = [DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice ([DimIndex SubExp] -> Slice SubExp)
-> [DimIndex SubExp] -> Slice SubExp
forall a b. (a -> b) -> a -> b
$ (SubExp -> DimIndex SubExp) -> [SubExp] -> [DimIndex SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimIndex SubExp
sliceDim [SubExp]
dims [DimIndex SubExp] -> [DimIndex SubExp] -> [DimIndex SubExp]
forall a. Semigroup a => a -> a -> a
<> Slice SubExp -> [DimIndex SubExp]
forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slice
      addStm $ Let pat_tan aux $ BasicOp $ Index arr_tan slice'
    Update Safety
safety VName
arr Slice SubExp
slice SubExp
se -> do
      dims <- Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims (Shape -> [SubExp]) -> ADM Shape -> ADM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ADM Shape
askShape
      arr_tan <- tangent arr
      se_tan <- tangent se
      let slice' = [DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice ([DimIndex SubExp] -> Slice SubExp)
-> [DimIndex SubExp] -> Slice SubExp
forall a b. (a -> b) -> a -> b
$ (SubExp -> DimIndex SubExp) -> [SubExp] -> [DimIndex SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimIndex SubExp
sliceDim [SubExp]
dims [DimIndex SubExp] -> [DimIndex SubExp] -> [DimIndex SubExp]
forall a. Semigroup a => a -> a -> a
<> Slice SubExp -> [DimIndex SubExp]
forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slice
      addStm $ Let pat_tan aux $ BasicOp $ Update safety arr_tan slice' se_tan
    Concat Int
d (VName
arr :| [VName]
arrs) SubExp
w -> do
      r <- Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (Shape -> Int) -> ADM Shape -> ADM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ADM Shape
askShape
      arr_tan <- tangent arr
      arrs_tans <- mapM tangent arrs
      addStm $ Let pat_tan aux $ BasicOp $ Concat (d + r) (arr_tan :| arrs_tans) w
    Manifest VName
arr [Int]
ds -> do
      r <- Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (Shape -> Int) -> ADM Shape -> ADM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ADM Shape
askShape
      arr_tan <- tangent arr
      addStm . Let pat_tan aux . BasicOp $
        Manifest arr_tan ([0 .. r - 1] ++ map (+ r) ds)
    Iota SubExp
n SubExp
_ SubExp
_ IntType
it -> do
      shape <- ADM Shape
askShape
      addStm . Let pat_tan aux . BasicOp $
        Replicate (shape <> Shape [n]) (intConst it 0)
    Replicate Shape
n SubExp
x ->
      Pat Type
-> StmAux () -> SubExp -> (SubExp -> ADM (Exp SOACS)) -> ADM ()
bindTan Pat Type
pat_tan StmAux ()
aux SubExp
x ((SubExp -> ADM (Exp SOACS)) -> ADM ())
-> (SubExp -> ADM (Exp SOACS)) -> ADM ()
forall a b. (a -> b) -> a -> b
$ \SubExp
x_tan ->
        Exp SOACS -> ADM (Exp SOACS)
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp SOACS -> ADM (Exp SOACS)) -> Exp SOACS -> ADM (Exp SOACS)
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ Shape -> SubExp -> BasicOp
Replicate Shape
n SubExp
x_tan
    Scratch PrimType
t [SubExp]
shape -> do
      tan_shape <- ADM Shape
askShape
      addStm $ Let pat_tan aux $ BasicOp $ Scratch t $ shapeDims tan_shape <> shape
    Reshape VName
arr NewShape SubExp
reshape -> do
      shape <- ADM Shape
askShape
      arr_tan <- tangent arr
      addStm $ Let pat_tan aux $ BasicOp $ Reshape arr_tan (newshapeInner shape reshape)
    Rearrange VName
arr [Int]
perm -> do
      r <- Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (Shape -> Int) -> ADM Shape -> ADM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ADM Shape
askShape
      arr_tan <- tangent arr
      addStm . Let pat_tan aux . BasicOp $
        Rearrange arr_tan ([0 .. r - 1] <> map (+ r) perm)
    BasicOp
_ -> [Char] -> ADM ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> ADM ()) -> [Char] -> ADM ()
forall a b. (a -> b) -> a -> b
$ [Char]
"basicFwd: Unsupported op " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ BasicOp -> [Char]
forall a. Pretty a => a -> [Char]
prettyString BasicOp
op

fwdLambda :: Lambda SOACS -> ADM (Lambda SOACS)
fwdLambda :: Lambda SOACS -> ADM (Lambda SOACS)
fwdLambda (Lambda [LParam SOACS]
params [Type]
_ Body SOACS
body) = do
  params' <- [Param Type] -> ADM [Param Type]
forall a. TanBuilder a => [a] -> ADM [a]
bundleNewList [Param Type]
[LParam SOACS]
params
  mkLambda params' $ bodyBind =<< fwdBody body

-- Differentiating a FlatMap lambda is mostly straightforward, except that we do
-- not care about the tangent of the size result. We must also take care to
-- reconstruct the return type, since we do not have facilities similar to
-- mkLambda for ExtLambdas. Finally, in the vector case (see Note [Forward-Mode
-- vector AD]) the tangent of a nonuniform result has the vector shape
-- outermost, but a FlatMap demands that the existential size come first, so we
-- must transpose it. It is transposed back outside the FlatMap; see 'fwdSOAC'.
fwdFlatMapLambda :: ExtLambda SOACS -> ADM (ExtLambda SOACS)
fwdFlatMapLambda :: ExtLambda SOACS -> ADM (ExtLambda SOACS)
fwdFlatMapLambda (Lambda [LParam SOACS]
params [ExtType]
rettype Body SOACS
body) = do
  params' <- [Param Type] -> ADM [Param Type]
forall a. TanBuilder a => [a] -> ADM [a]
bundleNewList [Param Type]
[LParam SOACS]
params
  (body', rettype') <- buildBody . localScope (scopeOfLParams params') $ do
    (meta_res, val_res) <- fmap (splitAt 2) <$> bodyBind =<< fwdBody body
    let (val_res_primal, val_res_tan) = unterleave val_res
        val_rettype = Int -> [ExtType] -> [ExtType]
forall a. Int -> [a] -> [a]
drop Int
1 [ExtType]
rettype
    val_rettype_tan <- mapM tanType val_rettype
    (val_res_tan', val_rettype_tan') <-
      mapAndUnzipM pullExtDim $ zip3 val_rettype val_res_tan val_rettype_tan
    pure
      ( take 1 meta_res <> interleave val_res_primal val_res_tan',
        Prim int64 : interleave val_rettype val_rettype_tan'
      )
  pure $ Lambda params' rettype' body'
  where
    -- Put the existential size back in the outermost position.
    pullExtDim :: (ExtType, SubExpRes, TypeBase (ShapeBase d) m)
-> ADM (SubExpRes, TypeBase (ShapeBase d) m)
pullExtDim (ExtType
t, res :: SubExpRes
res@(SubExpRes Certs
cs SubExp
se), TypeBase (ShapeBase d) m
t_tan)
      | ExtType -> Bool
flatMapNonuniform ExtType
t = do
          tan_shape <- ADM Shape
askShape
          v <- pushTanShape =<< asVName se
          pure
            ( SubExpRes cs $ Var v,
              rearrangeType (vecPerm tan_shape t_tan) t_tan
            )
      | Bool
otherwise = (SubExpRes, TypeBase (ShapeBase d) m)
-> ADM (SubExpRes, TypeBase (ShapeBase d) m)
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExpRes
res, TypeBase (ShapeBase d) m
t_tan)

fwdWithAccLambda :: [WithAccInput SOACS] -> Lambda SOACS -> ADM (Lambda SOACS)
fwdWithAccLambda :: [WithAccInput SOACS] -> Lambda SOACS -> ADM (Lambda SOACS)
fwdWithAccLambda [WithAccInput SOACS]
inputs (Lambda [LParam SOACS]
params [Type]
_ Body SOACS
body) = do
  let ([Param Type]
cert_params, [Param Type]
acc_params) = Int -> [Param Type] -> ([Param Type], [Param Type])
forall a. Int -> [a] -> ([a], [a])
splitAt ([WithAccInput SOACS] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [WithAccInput SOACS]
inputs) [Param Type]
[LParam SOACS]
params
  cert_params_tan <- Int -> ADM (Param Type) -> ADM [Param Type]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM ([WithAccInput SOACS] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [WithAccInput SOACS]
inputs) (ADM (Param Type) -> ADM [Param Type])
-> ADM (Param Type) -> ADM [Param Type]
forall a b. (a -> b) -> a -> b
$ Name -> Type -> ADM (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"acc_cert_tan" (Type -> ADM (Param Type)) -> Type -> ADM (Param Type)
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Unit
  acc_params_tan <- zipWithM mkAccParam (map paramName cert_params_tan) inputs

  mkLambda (cert_params <> cert_params_tan <> acc_params <> acc_params_tan) $ do
    zipWithM_
      insertTan
      (map paramName (cert_params <> acc_params))
      (map paramName (cert_params_tan <> acc_params_tan))
    bodyBind =<< fwdBody body
  where
    mkAccParam :: VName -> (Shape, [VName], c) -> ADM (Param (TypeBase shape o))
mkAccParam VName
c (Shape
shape, [VName]
arrs, c
_) = do
      tan_shape <- ADM Shape
askShape
      ts <- map (stripArray (shapeRank shape)) <$> mapM lookupType arrs
      newParam "acc_p_tan" $ Acc c (tan_shape <> shape) ts

fwdStreamLambda :: Int -> Lambda SOACS -> ADM (Lambda SOACS)
fwdStreamLambda :: Int -> Lambda SOACS -> ADM (Lambda SOACS)
fwdStreamLambda Int
num_accs (Lambda [LParam SOACS]
params [Type]
_ Body SOACS
body) = do
  tan_shape <- ADM Shape
askShape
  let (chunk_params, acc_params, arr_params) = splitAt3 1 num_accs params
  acc_params' <- bundleNewList acc_params
  (arr_params', arr_params'_tan) <- mapAndUnzipM onArrParam arr_params
  let params' =
        [Param Type]
chunk_params [Param Type] -> [Param Type] -> [Param Type]
forall a. Semigroup a => a -> a -> a
<> [Param Type]
acc_params' [Param Type] -> [Param Type] -> [Param Type]
forall a. Semigroup a => a -> a -> a
<> [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
interleave [Param Type]
arr_params' [Param Type]
arr_params'_tan
  mkLambda params' $ do
    zipWithM_ (trArrParamTan tan_shape) arr_params' arr_params'_tan
    (acc_res, map_res) <- fmap (splitAt (num_accs * 2)) . bodyBind =<< fwdBody body
    let (map_res_primal, map_res_tan) = unterleave map_res
    map_res_tan' <- mapM trMapResTan map_res_tan
    pure $ acc_res <> interleave map_res_primal map_res_tan'
  where
    -- Array parameters need to be treated specially as the chunk parameter
    -- must always be outermost.
    onArrParam :: Param dec -> ADM (Param dec, Param Type)
onArrParam Param dec
p = do
      shape <- ADM Shape
askShape
      (p', p_tan) <- bundleNew p
      let perm = Shape -> Type -> [Int]
forall s o. ArrayShape s => Shape -> TypeBase s o -> [Int]
vecPerm Shape
shape (Type -> [Int]) -> Type -> [Int]
forall a b. (a -> b) -> a -> b
$ Param dec -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param dec
p_tan
      pure (p', p_tan {paramDec = rearrangeType perm (paramType p_tan)})

    -- Put the tangent shape back in the outermost position.
    trArrParamTan :: Shape -> Param dec -> Param dec -> ADM ()
trArrParamTan Shape
tan_shape Param dec
p Param dec
p_tan = do
      let perm :: [Int]
perm = [Int] -> [Int]
rearrangeInverse ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ Shape -> Type -> [Int]
forall s o. ArrayShape s => Shape -> TypeBase s o -> [Int]
vecPerm Shape
tan_shape (Type -> [Int]) -> Type -> [Int]
forall a b. (a -> b) -> a -> b
$ Param dec -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param dec
p_tan
      v <-
        Name -> Exp (Rep ADM) -> ADM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp (VName -> Name
baseName (Param dec -> VName
forall dec. Param dec -> VName
paramName Param dec
p_tan)) (Exp SOACS -> ADM VName)
-> (BasicOp -> Exp SOACS) -> BasicOp -> ADM VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> ADM VName) -> BasicOp -> ADM VName
forall a b. (a -> b) -> a -> b
$
          VName -> [Int] -> BasicOp
Rearrange (Param dec -> VName
forall dec. Param dec -> VName
paramName Param dec
p_tan) [Int]
perm
      insertTan (paramName p) v

    -- Put the chunk size back in the outermost position.
    trMapResTan :: SubExpRes -> ADM SubExpRes
trMapResTan (SubExpRes Certs
cs ~(Var VName
v)) =
      Certs -> SubExp -> SubExpRes
SubExpRes Certs
cs (SubExp -> SubExpRes) -> (VName -> SubExp) -> VName -> SubExpRes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var (VName -> SubExpRes) -> ADM VName -> ADM SubExpRes
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> ADM VName
pushTanShape VName
v

pushTanShape :: VName -> ADM VName
pushTanShape :: VName -> ADM VName
pushTanShape VName
v = do
  tan_shape <- ADM Shape
askShape
  v_t <- lookupType v
  if tan_shape == mempty || arrayShape v_t == tan_shape || isAcc v_t
    then pure v
    else do
      let perm = Shape -> Type -> [Int]
forall s o. ArrayShape s => Shape -> TypeBase s o -> [Int]
vecPerm Shape
tan_shape Type
v_t
      letExp (baseName v <> "_tr") $ BasicOp $ Rearrange v perm

soacInputsWithTangents :: [VName] -> ADM [VName]
soacInputsWithTangents :: [VName] -> ADM [VName]
soacInputsWithTangents [VName]
xs = do
  xs_tans <- (VName -> ADM VName) -> [VName] -> ADM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (VName -> ADM VName
pushTanShape (VName -> ADM VName) -> (VName -> ADM VName) -> VName -> ADM VName
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< VName -> ADM VName
forall a. Tangent a => a -> ADM a
tangent) [VName]
xs
  pure $ interleave xs xs_tans

soacResPat :: Int -> Int -> Pat Type -> ADM (Pat Type, [(Pat Type, VName)])
soacResPat :: Int -> Int -> Pat Type -> ADM (Pat Type, [(Pat Type, VName)])
soacResPat Int
scan_res Int
red_res (Pat [PatElem Type]
pes) = do
  pes_tan <- (PatElem Type -> ADM (PatElem Type))
-> [PatElem Type] -> ADM [PatElem Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM PatElem Type -> ADM (PatElem Type)
forall a. TanBuilder a => a -> ADM a
newTan [PatElem Type]
pes
  bimap (Pat . interleave pes) mconcat . unzip <$> zipWithM tweakPatElem [0 ..] pes_tan
  where
    isRedRes :: Int -> Bool
isRedRes Int
i = Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
scan_res Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
scan_res Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
red_res
    tweakPatElem :: Int -> PatElem Type -> ADM (PatElem Type, [(Pat Type, VName)])
tweakPatElem Int
i pe :: PatElem Type
pe@(PatElem VName
v Type
v_t) = do
      tan_shape <- ADM Shape
askShape
      if isRedRes i || tan_shape == mempty || arrayShape v_t == tan_shape || isAcc v_t
        then pure (pe, [])
        else do
          let perm = Shape -> Type -> [Int]
forall s o. ArrayShape s => Shape -> TypeBase s o -> [Int]
vecPerm Shape
tan_shape Type
v_t
          v' <- newName v
          pure (PatElem v' $ rearrangeType perm v_t, [(Pat [pe], v')])

fwdSOAC :: Pat Type -> StmAux () -> SOAC SOACS -> ADM ()
fwdSOAC :: Pat Type -> StmAux () -> SOAC SOACS -> ADM ()
fwdSOAC Pat Type
pat StmAux ()
aux (Screma SubExp
size [VName]
xs (ScremaForm Lambda SOACS
f [Scan SOACS]
scs [Reduce SOACS]
reds Lambda SOACS
post_lam)) = do
  (pat', to_transpose) <- Int -> Int -> Pat Type -> ADM (Pat Type, [(Pat Type, VName)])
soacResPat ([Scan SOACS] -> Int
forall rep. [Scan rep] -> Int
scanResults [Scan SOACS]
scs) ([Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce SOACS]
reds) Pat Type
pat
  xs' <- soacInputsWithTangents xs
  f' <- fwdLambda f
  scs' <- mapM fwdScan scs
  reds' <- mapM fwdRed reds
  post_lam' <- fwdLambda post_lam
  addStm $ Let pat' aux $ Op $ Screma size xs' $ ScremaForm f' scs' reds' post_lam'
  tan_shape <- askShape
  forM_ to_transpose $ \(Pat Type
rpat, VName
v) -> do
    v_t <- VName -> ADM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
    let perm = [Int] -> [Int]
rearrangeInverse ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ Shape -> Type -> [Int]
forall s o. ArrayShape s => Shape -> TypeBase s o -> [Int]
vecPerm Shape
tan_shape Type
v_t
    letBind rpat $ BasicOp $ Rearrange v perm
  where
    zeroTans :: GLambda rep Type -> ADM [SubExp]
zeroTans GLambda rep Type
lam =
      (Type -> ADM SubExp) -> [Type] -> ADM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Name -> Exp (Rep ADM) -> ADM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"zero" (Exp SOACS -> ADM SubExp)
-> (Type -> Exp SOACS) -> Type -> ADM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Exp SOACS
zeroExp (Type -> ADM SubExp) -> (Type -> ADM Type) -> Type -> ADM SubExp
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Type -> ADM Type
forall s o.
(FromShape s, Monoid o) =>
TypeBase s o -> ADM (TypeBase s o)
tanType) ([Type] -> ADM [SubExp]) -> [Type] -> ADM [SubExp]
forall a b. (a -> b) -> a -> b
$ GLambda rep Type -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType GLambda rep Type
lam

    fwdScan :: Scan SOACS -> ADM (Scan SOACS)
    fwdScan :: Scan SOACS -> ADM (Scan SOACS)
fwdScan Scan SOACS
sc = do
      op' <- Lambda SOACS -> ADM (Lambda SOACS)
fwdLambda (Lambda SOACS -> ADM (Lambda SOACS))
-> Lambda SOACS -> ADM (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ Scan SOACS -> Lambda SOACS
forall rep. Scan rep -> Lambda rep
scanLambda Scan SOACS
sc
      neutral_tans <- zeroTans $ scanLambda sc
      pure $
        Scan
          { scanNeutral = scanNeutral sc `interleave` neutral_tans,
            scanLambda = op'
          }
    fwdRed :: Reduce SOACS -> ADM (Reduce SOACS)
    fwdRed :: Reduce SOACS -> ADM (Reduce SOACS)
fwdRed Reduce SOACS
red = do
      op' <- Lambda SOACS -> ADM (Lambda SOACS)
fwdLambda (Lambda SOACS -> ADM (Lambda SOACS))
-> Lambda SOACS -> ADM (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ Reduce SOACS -> Lambda SOACS
forall rep. Reduce rep -> Lambda rep
redLambda Reduce SOACS
red
      neutral_tans <- zeroTans $ redLambda red
      pure $
        Reduce
          { redComm = redComm red,
            redLambda = op',
            redNeutral = redNeutral red `interleave` neutral_tans
          }
fwdSOAC Pat Type
pat StmAux ()
aux (Stream SubExp
size [VName]
xs [SubExp]
accs Lambda SOACS
lam) = do
  pat' <- Pat Type -> ADM (Pat Type)
forall t. TanBuilder (PatElem t) => Pat t -> ADM (Pat t)
bundleNewPat Pat Type
pat
  lam' <- fwdStreamLambda (length accs) lam
  xs' <- soacInputsWithTangents xs
  accs_tan <- mapM (letSubExp "zero" . zeroExp <=< tanType <=< subExpType) accs
  let accs' = [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
interleave [SubExp]
accs [SubExp]
accs_tan
  addStm $ Let pat' aux $ Op $ Stream size xs' accs' lam'
fwdSOAC Pat Type
pat StmAux ()
aux (Hist SubExp
w [VName]
arrs [HistOp SOACS]
ops Lambda SOACS
bucket_fun) = do
  -- TODO: this is probably not very efficient in the vector case as we end up
  -- with a dreadful update operator that involves arrays.
  (pat', to_transpose) <- Int -> Int -> Pat Type -> ADM (Pat Type, [(Pat Type, VName)])
soacResPat Int
0 Int
0 Pat Type
pat
  ops' <- mapM fwdHist ops
  bucket_fun' <- fwdHistBucket bucket_fun
  arrs' <- soacInputsWithTangents arrs
  addStm $ Let pat' aux $ Op $ Hist w arrs' ops' bucket_fun'
  tan_shape <- askShape
  forM_ to_transpose $ \(Pat Type
rpat, VName
v) -> do
    v_t <- VName -> ADM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
    let perm = [Int] -> [Int]
rearrangeInverse ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ Shape -> Type -> [Int]
forall s o. ArrayShape s => Shape -> TypeBase s o -> [Int]
vecPerm Shape
tan_shape Type
v_t
    letBind rpat $ BasicOp $ Rearrange v perm
  where
    n_indices :: Int
n_indices = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (HistOp SOACS -> Int) -> [HistOp SOACS] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (Shape -> Int) -> (HistOp SOACS -> Shape) -> HistOp SOACS -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HistOp SOACS -> Shape
forall rep. HistOp rep -> Shape
histShape) [HistOp SOACS]
ops
    fwdBodyHist :: Body SOACS -> ADM (Body (Rep ADM))
fwdBodyHist (Body BodyDec SOACS
_ Stms SOACS
stms Result
res) = ADM Result -> ADM (Body (Rep ADM))
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (ADM Result -> ADM (Body (Rep ADM)))
-> ADM Result -> ADM (Body (Rep ADM))
forall a b. (a -> b) -> a -> b
$ do
      (Stm SOACS -> ADM ()) -> Stms SOACS -> ADM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Stm SOACS -> ADM ()
fwdStm Stms SOACS
stms
      let (Result
res_is, Result
res_vs) = Int -> Result -> (Result, Result)
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n_indices Result
res
      (Result
res_is Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++) (Result -> Result) -> ADM Result -> ADM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Result -> ADM Result
forall a. Tangent a => [a] -> ADM [a]
bundleTangents Result
res_vs
    fwdHistBucket :: Lambda SOACS -> ADM (Lambda (Rep ADM))
fwdHistBucket (Lambda [LParam SOACS]
params [Type]
_ Body SOACS
body) = do
      params' <- [Param Type] -> ADM [Param Type]
forall a. TanBuilder a => [a] -> ADM [a]
bundleNewList [Param Type]
[LParam SOACS]
params
      mkLambda params' $ bodyBind =<< fwdBodyHist body

    fwdHist :: HistOp SOACS -> ADM (HistOp SOACS)
    fwdHist :: HistOp SOACS -> ADM (HistOp SOACS)
fwdHist (HistOp Shape
shape SubExp
rf [VName]
dest [SubExp]
nes Lambda SOACS
op) = do
      dest' <- [VName] -> ADM [VName]
soacInputsWithTangents [VName]
dest
      nes_tan <- mapM (letSubExp "zero" . zeroExp <=< tanType) $ lambdaReturnType op
      op' <- fwdLambda op
      pure $
        HistOp
          { histShape = shape,
            histRaceFactor = rf,
            histDest = dest',
            histNeutral = interleave nes nes_tan,
            histOp = op'
          }
fwdSOAC Pat Type
pat StmAux ()
aux (WithVJP [SubExp]
args Lambda SOACS
lam Lambda SOACS
_) = do
  -- You have a custom adjoint? Too bad we are in tangent land.
  ((Stm SOACS -> ADM ()) -> Stms SOACS -> ADM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Stm SOACS -> ADM ()
fwdStm (Stms SOACS -> ADM ())
-> (Builder SOACS [()] -> ADM (Stms SOACS))
-> Builder SOACS [()]
-> ADM ()
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Builder SOACS [()] -> ADM (Stms SOACS)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (Stms rep)
runBuilder_) (Builder SOACS [()] -> ADM ()) -> Builder SOACS [()] -> ADM ()
forall a b. (a -> b) -> a -> b
$ do
    lam_res <- StmAux ()
-> BuilderT SOACS (State VNameSource) Result
-> BuilderT SOACS (State VNameSource) Result
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux ()
aux (BuilderT SOACS (State VNameSource) Result
 -> BuilderT SOACS (State VNameSource) Result)
-> BuilderT SOACS (State VNameSource) Result
-> BuilderT SOACS (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$ Lambda (Rep (BuilderT SOACS (State VNameSource)))
-> [BuilderT
      SOACS
      (State VNameSource)
      (Exp (Rep (BuilderT SOACS (State VNameSource))))]
-> BuilderT SOACS (State VNameSource) Result
forall (m :: * -> *).
MonadBuilder m =>
Lambda (Rep m) -> [m (Exp (Rep m))] -> m Result
eLambda Lambda (Rep (BuilderT SOACS (State VNameSource)))
Lambda SOACS
lam ([BuilderT
    SOACS
    (State VNameSource)
    (Exp (Rep (BuilderT SOACS (State VNameSource))))]
 -> BuilderT SOACS (State VNameSource) Result)
-> [BuilderT
      SOACS
      (State VNameSource)
      (Exp (Rep (BuilderT SOACS (State VNameSource))))]
-> BuilderT SOACS (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$ (SubExp
 -> BuilderT
      SOACS
      (State VNameSource)
      (Exp (Rep (BuilderT SOACS (State VNameSource)))))
-> [SubExp]
-> [BuilderT
      SOACS
      (State VNameSource)
      (Exp (Rep (BuilderT SOACS (State VNameSource))))]
forall a b. (a -> b) -> [a] -> [b]
map SubExp
-> BuilderT
     SOACS
     (State VNameSource)
     (Exp (Rep (BuilderT SOACS (State VNameSource))))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp [SubExp]
args
    forM (zip (patNames pat) lam_res) $ \(VName
v, SubExpRes Certs
cs SubExp
se) ->
      Certs
-> BuilderT SOACS (State VNameSource) ()
-> BuilderT SOACS (State VNameSource) ()
forall a.
Certs
-> BuilderT SOACS (State VNameSource) a
-> BuilderT SOACS (State VNameSource) a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (BuilderT SOACS (State VNameSource) ()
 -> BuilderT SOACS (State VNameSource) ())
-> BuilderT SOACS (State VNameSource) ()
-> BuilderT SOACS (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ [VName]
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v] (Exp (Rep (BuilderT SOACS (State VNameSource)))
 -> BuilderT SOACS (State VNameSource) ())
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource))))
-> BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
fwdSOAC (Pat [PatElem Type]
pes) StmAux ()
aux (FlatMap SubExp
w [VName]
arrs ExtLambda SOACS
lam) = do
  -- We do not bother to create tangents for the metadata part of the result.
  let ([PatElem Type]
meta_pes, [PatElem Type]
val_pes) = Int -> [PatElem Type] -> ([PatElem Type], [PatElem Type])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
4 [PatElem Type]
pes
  (Pat val_pes', to_transpose) <- Int -> Int -> Pat Type -> ADM (Pat Type, [(Pat Type, VName)])
soacResPat Int
0 Int
0 (Pat Type -> ADM (Pat Type, [(Pat Type, VName)]))
-> Pat Type -> ADM (Pat Type, [(Pat Type, VName)])
forall a b. (a -> b) -> a -> b
$ [PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat [PatElem Type]
val_pes
  let pat' = [PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat ([PatElem Type] -> Pat Type) -> [PatElem Type] -> Pat Type
forall a b. (a -> b) -> a -> b
$ [PatElem Type]
meta_pes [PatElem Type] -> [PatElem Type] -> [PatElem Type]
forall a. Semigroup a => a -> a -> a
<> [PatElem Type]
val_pes'
  arrs' <- soacInputsWithTangents arrs
  lam' <- fwdFlatMapLambda lam
  addStm $ Let pat' aux $ Op $ FlatMap w arrs' lam'
  tan_shape <- askShape
  forM_ to_transpose $ \(Pat Type
rpat, VName
v) -> do
    v_t <- VName -> ADM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
    let perm = [Int] -> [Int]
rearrangeInverse ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ Shape -> Type -> [Int]
forall s o. ArrayShape s => Shape -> TypeBase s o -> [Int]
vecPerm Shape
tan_shape Type
v_t
    letBind rpat $ BasicOp $ Rearrange v perm
fwdSOAC Pat Type
_ StmAux ()
_ JVP {} =
  [Char] -> ADM ()
forall a. HasCallStack => [Char] -> a
error [Char]
"fwdSOAC: nested JVP not allowed."
fwdSOAC Pat Type
_ StmAux ()
_ VJP {} =
  [Char] -> ADM ()
forall a. HasCallStack => [Char] -> a
error [Char]
"fwdSOAC: nested VJP not allowed."

fwdStm :: Stm SOACS -> ADM ()
fwdStm :: Stm SOACS -> ADM ()
fwdStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (BasicOp (UpdateAcc Safety
safety VName
acc [SubExp]
i [SubExp]
x))) = do
  pat_tan <- Pat Type -> ADM (Pat Type)
forall t. TanBuilder (PatElem t) => Pat t -> ADM (Pat t)
newTanPat Pat Type
Pat (LetDec SOACS)
pat
  addStm $ Let pat aux $ BasicOp $ UpdateAcc safety acc i x
  addStm . Let pat_tan aux <=< withTansI acc x $ \[SubExp]
is VName
acc_tan [SubExp]
x_tan' -> do
    Exp SOACS -> ADM (Exp SOACS)
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp SOACS -> ADM (Exp SOACS)) -> Exp SOACS -> ADM (Exp SOACS)
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ Safety -> VName -> [SubExp] -> [SubExp] -> BasicOp
UpdateAcc Safety
safety VName
acc_tan ([SubExp]
is [SubExp] -> [SubExp] -> [SubExp]
forall a. Semigroup a => a -> a -> a
<> [SubExp]
i) [SubExp]
x_tan'
fwdStm stm :: Stm SOACS
stm@(Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (BasicOp BasicOp
e)) = do
  -- XXX: this has to be too naive.
  Bool -> ADM () -> ADM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ((Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc ([Type] -> Bool) -> [Type] -> Bool
forall a b. (a -> b) -> a -> b
$ Pat Type -> [Type]
forall dec. Typed dec => Pat dec -> [Type]
patTypes Pat Type
Pat (LetDec SOACS)
pat) (ADM () -> ADM ()) -> ADM () -> ADM ()
forall a b. (a -> b) -> a -> b
$ Stm (Rep ADM) -> ADM ()
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm Stm (Rep ADM)
Stm SOACS
stm
  Pat Type -> StmAux () -> BasicOp -> ADM ()
basicFwd Pat Type
Pat (LetDec SOACS)
pat StmAux ()
StmAux (ExpDec SOACS)
aux BasicOp
e
fwdStm stm :: Stm SOACS
stm@(Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Apply Name
f [(SubExp, Diet)]
args [(RetType SOACS, RetAls)]
_ Safety
_))
  | Just (PrimType
ret, [PrimType]
argts) <- Name
-> Map Name (PrimType, [PrimType]) -> Maybe (PrimType, [PrimType])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name (PrimType, [PrimType])
builtInFunctions = do
      Stm (Rep ADM) -> ADM ()
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm Stm (Rep ADM)
Stm SOACS
stm
      pat_tan <- Pat Type -> ADM (Pat Type)
forall t. TanBuilder (PatElem t) => Pat t -> ADM (Pat t)
newTanPat Pat Type
Pat (LetDec SOACS)
pat
      let arg_pes = (PrimType -> SubExp -> PrimExp VName)
-> [PrimType] -> [SubExp] -> [PrimExp VName]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PrimType -> SubExp -> PrimExp VName
primExpFromSubExp [PrimType]
argts (((SubExp, Diet) -> SubExp) -> [(SubExp, Diet)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp, Diet) -> SubExp
forall a b. (a, b) -> a
fst [(SubExp, Diet)]
args)
      case pdBuiltin f arg_pes of
        Maybe [PrimExp VName]
Nothing ->
          [Char] -> ADM ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> ADM ()) -> [Char] -> ADM ()
forall a b. (a -> b) -> a -> b
$ [Char]
"No partial derivative defined for builtin function: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Name -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Name
f
        Just [PrimExp VName]
derivs -> do
          let convertTo :: PrimType -> PrimExp VName -> PrimExp VName
convertTo PrimType
tt PrimExp VName
e
                | PrimType
e_t PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType
tt = PrimExp VName
e
                | Bool
otherwise =
                    case (PrimType
tt, PrimType
e_t) of
                      (IntType IntType
tt', IntType IntType
ft) -> ConvOp -> PrimExp VName -> PrimExp VName
forall v. ConvOp -> PrimExp v -> PrimExp v
ConvOpExp (IntType -> IntType -> ConvOp
SExt IntType
ft IntType
tt') PrimExp VName
e
                      (FloatType FloatType
tt', FloatType FloatType
ft) -> ConvOp -> PrimExp VName -> PrimExp VName
forall v. ConvOp -> PrimExp v -> PrimExp v
ConvOpExp (FloatType -> FloatType -> ConvOp
FPConv FloatType
ft FloatType
tt') PrimExp VName
e
                      (PrimType
Bool, FloatType FloatType
ft) -> ConvOp -> PrimExp VName -> PrimExp VName
forall v. ConvOp -> PrimExp v -> PrimExp v
ConvOpExp (FloatType -> ConvOp
FToB FloatType
ft) PrimExp VName
e
                      (FloatType FloatType
tt', PrimType
Bool) -> ConvOp -> PrimExp VName -> PrimExp VName
forall v. ConvOp -> PrimExp v -> PrimExp v
ConvOpExp (FloatType -> ConvOp
BToF FloatType
tt') PrimExp VName
e
                      (PrimType, PrimType)
_ -> [Char] -> PrimExp VName
forall a. HasCallStack => [Char] -> a
error ([Char] -> PrimExp VName) -> [Char] -> PrimExp VName
forall a b. (a -> b) -> a -> b
$ [Char]
"fwdStm.convertTo: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Name, PrimType, PrimType) -> [Char]
forall a. Pretty a => a -> [Char]
prettyString (Name
f, PrimType
tt, PrimType
e_t)
                where
                  e_t :: PrimType
e_t = PrimExp VName -> PrimType
forall v. PrimExp v -> PrimType
primExpType PrimExp VName
e

          Pat Type -> StmAux () -> Exp SOACS -> ADM ()
bindTanPat Pat Type
pat_tan StmAux ()
StmAux (ExpDec SOACS)
aux (Exp SOACS -> ADM ())
-> (([PrimExp VName] -> PrimExp VName) -> ADM (Exp SOACS))
-> ([PrimExp VName] -> PrimExp VName)
-> ADM ()
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< [SubExp] -> ([PrimExp VName] -> PrimExp VName) -> ADM (Exp SOACS)
forall (f :: * -> *).
Traversable f =>
f SubExp -> ([PrimExp VName] -> PrimExp VName) -> ADM (Exp SOACS)
withAnyTans (((SubExp, Diet) -> SubExp) -> [(SubExp, Diet)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp, Diet) -> SubExp
forall a b. (a, b) -> a
fst [(SubExp, Diet)]
args) (([PrimExp VName] -> PrimExp VName) -> ADM ())
-> ([PrimExp VName] -> PrimExp VName) -> ADM ()
forall a b. (a -> b) -> a -> b
$
            \[PrimExp VName]
arg_tans' ->
              (PrimExp VName -> PrimExp VName -> PrimExp VName)
-> [PrimExp VName] -> PrimExp VName
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. PrimExp v -> PrimExp v -> PrimExp v
(~+~) ([PrimExp VName] -> PrimExp VName)
-> [PrimExp VName] -> PrimExp VName
forall a b. (a -> b) -> a -> b
$ (PrimExp VName -> PrimExp VName -> PrimExp VName)
-> [PrimExp VName] -> [PrimExp VName] -> [PrimExp VName]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. PrimExp v -> PrimExp v -> PrimExp v
(~*~) ((PrimExp VName -> PrimExp VName)
-> [PrimExp VName] -> [PrimExp VName]
forall a b. (a -> b) -> [a] -> [b]
map (PrimType -> PrimExp VName -> PrimExp VName
convertTo PrimType
ret) [PrimExp VName]
arg_tans') [PrimExp VName]
derivs
fwdStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Match [SubExp]
ses [Case (Body SOACS)]
cases Body SOACS
defbody (MatchDec [BranchType SOACS]
ret MatchSort
ifsort))) = do
  cases' <- ADM [Case (Body SOACS)] -> ADM [Case (Body SOACS)]
forall a. ADM a -> ADM a
slocal' (ADM [Case (Body SOACS)] -> ADM [Case (Body SOACS)])
-> ADM [Case (Body SOACS)] -> ADM [Case (Body SOACS)]
forall a b. (a -> b) -> a -> b
$ (Case (Body SOACS) -> ADM (Case (Body SOACS)))
-> [Case (Body SOACS)] -> ADM [Case (Body SOACS)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((Body SOACS -> ADM (Body SOACS))
-> Case (Body SOACS) -> ADM (Case (Body SOACS))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Case a -> f (Case b)
traverse Body SOACS -> ADM (Body SOACS)
fwdBody) [Case (Body SOACS)]
cases
  defbody' <- slocal' $ fwdBody defbody
  pat' <- bundleNewPat pat
  ret' <- bundleTangents ret
  addStm $ Let pat' aux $ Match ses cases' defbody' $ MatchDec ret' ifsort
fwdStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Loop [(FParam SOACS, SubExp)]
val_pats loop :: LoopForm
loop@(WhileLoop VName
v) Body SOACS
body)) = do
  val_pats' <- [(Param DeclType, SubExp)] -> ADM [(Param DeclType, SubExp)]
forall a. TanBuilder a => [a] -> ADM [a]
bundleNewList [(Param DeclType, SubExp)]
[(FParam SOACS, SubExp)]
val_pats
  pat' <- bundleNewPat pat
  body' <-
    localScope (scopeOfFParams (map fst val_pats') <> scopeOfLoopForm loop) . slocal' $
      fwdBody body
  addStm $ Let pat' aux $ Loop val_pats' (WhileLoop v) body'
fwdStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Loop [(FParam SOACS, SubExp)]
val_pats loop :: LoopForm
loop@(ForLoop VName
i IntType
it SubExp
bound) Body SOACS
body)) = do
  pat' <- Pat Type -> ADM (Pat Type)
forall t. TanBuilder (PatElem t) => Pat t -> ADM (Pat t)
bundleNewPat Pat Type
Pat (LetDec SOACS)
pat
  val_pats' <- bundleNewList val_pats
  body' <-
    localScope (scopeOfFParams (map fst val_pats') <> scopeOfLoopForm loop) . slocal' $
      fwdBody body
  addStm $ Let pat' aux $ Loop val_pats' (ForLoop i it bound) body'
fwdStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (WithAcc [WithAccInput SOACS]
inputs Lambda SOACS
lam)) = do
  inputs_tan <- [WithAccInput SOACS]
-> (WithAccInput SOACS -> ADM (WithAccInput SOACS))
-> ADM [WithAccInput SOACS]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [WithAccInput SOACS]
inputs ((WithAccInput SOACS -> ADM (WithAccInput SOACS))
 -> ADM [WithAccInput SOACS])
-> (WithAccInput SOACS -> ADM (WithAccInput SOACS))
-> ADM [WithAccInput SOACS]
forall a b. (a -> b) -> a -> b
$ \(Shape
shape, [VName]
arrs, Maybe (Lambda SOACS, [SubExp])
op) -> do
    arrs_tan <- (VName -> ADM VName) -> [VName] -> ADM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> ADM VName
forall a. Tangent a => a -> ADM a
tangent [VName]
arrs
    tan_shape <- askShape
    op' <- case op of
      Maybe (Lambda SOACS, [SubExp])
Nothing -> Maybe (Lambda SOACS, [SubExp])
-> ADM (Maybe (Lambda SOACS, [SubExp]))
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Lambda SOACS, [SubExp])
forall a. Maybe a
Nothing
      Just (Lambda SOACS
op_lam, [SubExp]
nes) -> do
        -- We assume that op_lam has unit partial derivatives (i.e., is some
        -- kind of addition). This is the case for all WithAccs produced by VJP.
        -- The operator takes one index per dimension of the index space of the
        -- tangent accumulator, which includes the vector dimensions.
        op_lam' <-
          Int -> [Type] -> ADM (Lambda SOACS)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ SOACS) =>
Int -> [Type] -> m (Lambda SOACS)
accAddLambda
            (Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
tan_shape Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
shape)
            (Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
op_lam)
        pure $ Just (op_lam', nes)
    pure (tan_shape <> shape, arrs_tan, op')
  pat' <- bundleNewPat pat
  lam' <- fwdWithAccLambda inputs lam
  addStm $ Let pat' aux $ WithAcc (interleave inputs inputs_tan) lam'
fwdStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Op Op SOACS
soac)) = Pat Type -> StmAux () -> SOAC SOACS -> ADM ()
fwdSOAC Pat Type
Pat (LetDec SOACS)
pat StmAux ()
StmAux (ExpDec SOACS)
aux Op SOACS
SOAC SOACS
soac
fwdStm Stm SOACS
stm =
  [Char] -> ADM ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> ADM ()) -> [Char] -> ADM ()
forall a b. (a -> b) -> a -> b
$ [Char]
"unhandled forward mode AD for Stm: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Stm SOACS -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Stm SOACS
stm [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Stm SOACS -> [Char]
forall a. Show a => a -> [Char]
show Stm SOACS
stm

fwdBody :: Body SOACS -> ADM (Body SOACS)
fwdBody :: Body SOACS -> ADM (Body SOACS)
fwdBody (Body BodyDec SOACS
_ Stms SOACS
stms Result
res) = ADM Result -> ADM (Body (Rep ADM))
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (ADM Result -> ADM (Body (Rep ADM)))
-> ADM Result -> ADM (Body (Rep ADM))
forall a b. (a -> b) -> a -> b
$ do
  (Stm SOACS -> ADM ()) -> Stms SOACS -> ADM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Stm SOACS -> ADM ()
fwdStm Stms SOACS
stms
  Result -> ADM Result
forall a. Tangent a => [a] -> ADM [a]
bundleTangents Result
res

fwdBodyTansLast :: Body SOACS -> ADM (Body SOACS)
fwdBodyTansLast :: Body SOACS -> ADM (Body SOACS)
fwdBodyTansLast (Body BodyDec SOACS
_ Stms SOACS
stms Result
res) = ADM Result -> ADM (Body (Rep ADM))
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (ADM Result -> ADM (Body (Rep ADM)))
-> ADM Result -> ADM (Body (Rep ADM))
forall a b. (a -> b) -> a -> b
$ do
  (Stm SOACS -> ADM ()) -> Stms SOACS -> ADM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Stm SOACS -> ADM ()
fwdStm Stms SOACS
stms
  (Result
res Result -> Result -> Result
forall a. Semigroup a => a -> a -> a
<>) (Result -> Result) -> ADM Result -> ADM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExpRes -> ADM SubExpRes) -> Result -> ADM Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExpRes -> ADM SubExpRes
forall a. Tangent a => a -> ADM a
tangent Result
res

fwdJVP ::
  (MonadFreshNames m) =>
  Scope SOACS ->
  Shape ->
  Attrs ->
  Lambda SOACS ->
  m (Lambda SOACS)
fwdJVP :: forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Shape -> Attrs -> Lambda SOACS -> m (Lambda SOACS)
fwdJVP Scope SOACS
scope Shape
shape Attrs
attrs (Lambda [LParam SOACS]
params [Type]
_ Body SOACS
body) =
  Shape -> Attrs -> ADM (Lambda SOACS) -> m (Lambda SOACS)
forall (m :: * -> *) a.
MonadFreshNames m =>
Shape -> Attrs -> ADM a -> m a
runADM Shape
shape Attrs
attrs (ADM (Lambda SOACS) -> m (Lambda SOACS))
-> (ADM (Lambda SOACS) -> ADM (Lambda SOACS))
-> ADM (Lambda SOACS)
-> m (Lambda SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope SOACS -> ADM (Lambda SOACS) -> ADM (Lambda SOACS)
forall a. Scope SOACS -> ADM a -> ADM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope Scope SOACS
scope (ADM (Lambda SOACS) -> m (Lambda SOACS))
-> ADM (Lambda SOACS) -> m (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ do
    params_tan <- (Param Type -> ADM (Param Type))
-> [Param Type] -> ADM [Param Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Param Type -> ADM (Param Type)
forall a. TanBuilder a => a -> ADM a
newTan [Param Type]
[LParam SOACS]
params
    mkLambda (params <> params_tan) $
      bodyBind =<< fwdBodyTansLast body

-- Note [Forward-Mode vector AD]
--
-- An primal variable of type 't' has a tangent of type '[tan_shape]t', where
-- 'tan_shape' is the vector shape (which may be empty in the non-vector case).
-- This requires some care for SOACs, which always map across the outermost
-- dimension: basically we have to transpose the inputs and the outputs.