{-# LANGUAGE TypeFamilies #-}

-- Naming scheme:
--
-- An adjoint-related object for "x" is named "x_adj".  This means
-- both actual adjoints and statements.
--
-- Do not assume "x'" means anything related to derivatives.
module Futhark.AD.Rev (revVJP) where

import Control.Monad
import Data.List.NonEmpty (NonEmpty (..))
import Data.Map qualified as M
import Data.Tuple
import Futhark.AD.Derivatives
import Futhark.AD.Rev.Acc
import Futhark.AD.Rev.Loop
import Futhark.AD.Rev.Monad
import Futhark.AD.Rev.SOAC
import Futhark.AD.Shared
import Futhark.Analysis.PrimExp.Convert
import Futhark.Builder
import Futhark.IR.SOACS
import Futhark.Tools
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
import Futhark.Util (takeLast)

patName :: Pat Type -> ADM VName
patName :: Pat Type -> ADM VName
patName (Pat [PatElem Type
pe]) = VName -> ADM VName
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> ADM VName) -> VName -> ADM VName
forall a b. (a -> b) -> a -> b
$ PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
pe
patName Pat Type
pat = [Char] -> ADM VName
forall a. HasCallStack => [Char] -> a
error ([Char] -> ADM VName) -> [Char] -> ADM VName
forall a b. (a -> b) -> a -> b
$ [Char]
"Expected single-element pattern: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Pat Type -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Pat Type
pat

copyIfArray :: VName -> ADM VName
copyIfArray :: VName -> ADM VName
copyIfArray VName
v = do
  v_t <- VName -> ADM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
  case v_t of
    Array {} ->
      Name -> Exp (Rep ADM) -> ADM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp (VName -> Name
baseName VName
v Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_copy") (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
$ Shape -> SubExp -> BasicOp
Replicate Shape
forall a. Monoid a => a
mempty (VName -> SubExp
Var VName
v)
    Type
_ -> VName -> ADM VName
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v

-- The vast majority of BasicOps require no special treatment in the
-- forward pass and produce one value (and hence one adjoint).  We
-- deal with that case here.
commonBasicOp :: Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp :: Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
op ADM ()
m = do
  Stm (Rep ADM) -> ADM ()
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm (Stm (Rep ADM) -> ADM ()) -> Stm (Rep ADM) -> ADM ()
forall a b. (a -> b) -> a -> b
$ Pat (LetDec (Rep ADM))
-> StmAux (ExpDec (Rep ADM)) -> Exp (Rep ADM) -> Stm (Rep ADM)
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat Type
Pat (LetDec (Rep ADM))
pat StmAux ()
StmAux (ExpDec (Rep ADM))
aux (Exp (Rep ADM) -> Stm (Rep ADM)) -> Exp (Rep ADM) -> Stm (Rep ADM)
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep ADM)
forall rep. BasicOp -> Exp rep
BasicOp BasicOp
op
  ADM ()
m
  pat_v <- Pat Type -> ADM VName
patName Pat Type
pat
  pat_adj <- lookupAdjVal pat_v
  pure (pat_v, pat_adj)

diffBasicOp :: Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM ()
diffBasicOp :: Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM ()
diffBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m =
  case BasicOp
e of
    CmpOp {} ->
      ADM (VName, VName) -> ADM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ADM (VName, VName) -> ADM ()) -> ADM (VName, VName) -> ADM ()
forall a b. (a -> b) -> a -> b
$ Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
    --
    ConvOp ConvOp
op SubExp
x -> do
      (_pat_v, pat_adj) <- Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
      returnSweepCode $ do
        adj_shape <- askShape

        contrib <- letExp "convop_contrib" <=< mapNest adj_shape (MkSolo (Var pat_adj)) $
          \(MkSolo SubExp
pat_adj') ->
            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 -> ConvOp
flipConvOp ConvOp
op) SubExp
pat_adj'

        updateSubExpAdj x contrib
    --
    UnOp UnOp
op SubExp
x -> do
      (_pat_v, pat_adj) <- Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m

      returnSweepCode $ do
        let t = UnOp -> PrimType
unOpType UnOp
op

        adj_shape <- askShape

        contrib <- letExp "unop_contrib" <=< mapNest adj_shape (MkSolo (Var pat_adj)) $
          \(MkSolo SubExp
pat_adj') ->
            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
pat_adj' PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. PrimExp v -> PrimExp v -> PrimExp v
~*~ UnOp -> PrimExp VName -> PrimExp VName
pdUnOp UnOp
op (PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
t SubExp
x)

        updateSubExpAdj x contrib
    --
    BinOp BinOp
op SubExp
x SubExp
y -> do
      (_pat_v, pat_adj) <- Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m

      returnSweepCode $ do
        let t = BinOp -> PrimType
binOpType BinOp
op
            (wrt_x, wrt_y) =
              pdBinOp op (primExpFromSubExp t x) (primExpFromSubExp t y)

        adj_shape <- askShape

        adj_x <- letExp "binop_x_adj"
          <=< mapNest adj_shape (MkSolo (Var pat_adj))
          $ \(MkSolo SubExp
pat_adj') ->
            let pat_adj'' :: PrimExp VName
pat_adj'' = PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
t SubExp
pat_adj'
             in 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
$ PrimExp VName
pat_adj'' PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. PrimExp v -> PrimExp v -> PrimExp v
~*~ PrimExp VName
wrt_x

        adj_y <- letExp "binop_y_adj"
          <=< mapNest adj_shape (MkSolo (Var pat_adj))
          $ \(MkSolo SubExp
pat_adj') ->
            let pat_adj'' :: PrimExp VName
pat_adj'' = PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
t SubExp
pat_adj'
             in 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
$ PrimExp VName
pat_adj'' PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. PrimExp v -> PrimExp v -> PrimExp v
~*~ PrimExp VName
wrt_y

        updateSubExpAdj x adj_x
        updateSubExpAdj y adj_y
    --
    SubExp SubExp
se -> do
      (_pat_v, pat_adj) <- Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
      returnSweepCode $ updateSubExpAdj se pat_adj
    --
    Assert {} ->
      ADM (VName, VName) -> ADM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ADM (VName, VName) -> ADM ()) -> ADM (VName, VName) -> ADM ()
forall a b. (a -> b) -> a -> b
$ Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
    --
    ArrayVal {} ->
      ADM (VName, VName) -> ADM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ADM (VName, VName) -> ADM ()) -> ADM (VName, VName) -> ADM ()
forall a b. (a -> b) -> a -> b
$ Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
    --
    ArrayLit [SubExp]
elems Type
_ -> do
      (_pat_v, pat_adj) <- Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
      t <- lookupType pat_adj
      returnSweepCode $ do
        forM_ (zip [(0 :: Int64) ..] elems) $ \(Int64
i, SubExp
se) -> do
          slice <- Type -> [DimIndex SubExp] -> ADM (Slice SubExp)
vecSlice Type
t [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant Int64
i)]
          updateSubExpAdj se <=< letExp "elem_adj" $ BasicOp $ Index pat_adj slice
    --
    Index VName
arr Slice SubExp
slice -> do
      (_pat_v, pat_adj) <- Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
      returnSweepCode $ void $ updateAdjSlice slice arr pat_adj
    FlatIndex {} -> [Char] -> ADM ()
forall a. HasCallStack => [Char] -> a
error [Char]
"FlatIndex not handled by AD yet."
    FlatUpdate {} -> [Char] -> ADM ()
forall a. HasCallStack => [Char] -> a
error [Char]
"FlatUpdate not handled by AD yet."
    --
    Opaque OpaqueOp
_ SubExp
se -> do
      (_pat_v, pat_adj) <- Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
      returnSweepCode $ updateSubExpAdj se pat_adj
    --
    Reshape VName
arr NewShape SubExp
newshape -> do
      (_pat_v, pat_adj) <- Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
      returnSweepCode $ do
        adj_shape <- askShape
        arr_shape <- arrayShape <$> lookupType arr
        void $
          updateAdj arr <=< letExp "adj_reshape" . BasicOp $
            Reshape pat_adj . newshapeInner adj_shape $
              reshapeAll (newShape newshape) arr_shape
    --
    Rearrange VName
arr [Int]
perm -> do
      (_pat_v, pat_adj) <- Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
      r <- shapeRank <$> askShape
      returnSweepCode $
        void . updateAdj arr <=< letExp "adj_rearrange" . BasicOp $
          Rearrange pat_adj ([0 .. r - 1] <> map (+ r) (rearrangeInverse perm))
    --
    Replicate (Shape []) (Var VName
se) -> do
      (_pat_v, pat_adj) <- Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
      returnSweepCode $ void $ updateAdj se pat_adj
    --
    Replicate (Shape [SubExp]
ns) SubExp
x -> do
      (_pat_v, pat_adj) <- Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
      returnSweepCode $ do
        adj_shape <- askShape
        x_t <- subExpType x
        lam <- addLambda x_t
        ne <- letSubExp "zero" $ zeroExp x_t
        n <- letSubExp "rep_size" =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) ns
        reduce <- reduceSOAC [Reduce Commutative lam [ne]]

        contrib <- letExp "rep_contrib" <=< mapNest adj_shape (MkSolo (Var pat_adj)) $
          \(MkSolo SubExp
pat_adj') -> do
            pat_adj_v <- SubExp -> ADM VName
forall (m :: * -> *). MonadBuilder m => SubExp -> m VName
asVName SubExp
pat_adj'
            -- Flatten the replicated dimensions into a single dimension that we
            -- can reduce across.
            pat_adj_flat <-
              letExp (baseName pat_adj <> "_flat") . BasicOp . Reshape pat_adj_v $
                reshapeAll (Shape ns <> arrayShape x_t) (Shape $ n : arrayDims x_t)
            pure $ Op $ Screma n [pat_adj_flat] reduce

        updateSubExpAdj x contrib
    --
    Concat Int
d (VName
arr :| [VName]
arrs) SubExp
_ -> do
      (_pat_v, pat_adj) <- Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
      returnSweepCode $ do
        let sliceAdj SubExp
_ [] = [VName] -> ADM [VName]
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
            sliceAdj SubExp
start (VName
v : [VName]
vs) = do
              v_t <- VName -> ADM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
              pat_adj_t <- lookupType pat_adj
              r <- shapeRank <$> askShape
              let w = Int -> Type -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
d Type
v_t
                  slice = SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
DimSlice SubExp
start SubExp
w (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
              pat_adj_slice <-
                letExp (baseName pat_adj <> "_slice") $
                  BasicOp $
                    Index pat_adj (sliceAt pat_adj_t (r + d) [slice])
              start' <- letSubExp "start" $ BasicOp $ BinOp (Add Int64 OverflowUndef) start w
              slices <- sliceAdj start' vs
              pure $ pat_adj_slice : slices

        slices <- sliceAdj (intConst Int64 0) $ arr : arrs

        zipWithM_ updateAdj (arr : arrs) slices
    --
    Manifest VName
se [Int]
_ -> do
      (_pat_v, pat_adj) <- Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
      returnSweepCode $ void $ updateAdj se pat_adj
    --
    Scratch {} ->
      ADM (VName, VName) -> ADM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ADM (VName, VName) -> ADM ()) -> ADM (VName, VName) -> ADM ()
forall a b. (a -> b) -> a -> b
$ Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
    --
    Iota SubExp
n SubExp
_ SubExp
_ IntType
t -> do
      (_pat_v, pat_adj) <- Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
      returnSweepCode $ do
        ne <- letSubExp "zero" $ zeroExp $ Prim $ IntType t
        lam <- addLambda $ Prim $ IntType t
        reduce <- reduceSOAC [Reduce Commutative lam [ne]]
        updateSubExpAdj n
          =<< letExp "iota_contrib" (Op $ Screma n [pat_adj] reduce)
    --
    Update Safety
safety VName
arr Slice SubExp
slice SubExp
v -> do
      (_pat_v, pat_adj) <- Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m
      returnSweepCode $ do
        adj_shape <- askShape
        let adj_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 (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
adj_shape) [DimIndex SubExp] -> [DimIndex SubExp] -> [DimIndex SubExp]
forall a. [a] -> [a] -> [a]
++ Slice SubExp -> [DimIndex SubExp]
forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slice
        v_adj <- letExp "update_val_adj" $ BasicOp $ Index pat_adj adj_slice
        v_adj_copy <- copyIfArray v_adj
        updateSubExpAdj v v_adj_copy
        v_adj_t <- lookupType v_adj
        zeroes <- letSubExp "update_zero" $ zeroExp v_adj_t
        void $
          updateAdj arr
            =<< letExp "update_src_adj" (BasicOp $ Update safety pat_adj adj_slice zeroes)
    UpdateAcc Safety
safety VName
acc [SubExp]
is [SubExp]
vs ->
      Pat Type
-> StmAux ()
-> Safety
-> VName
-> [SubExp]
-> [SubExp]
-> ADM ()
-> ADM ()
diffUpdateAcc Pat Type
pat StmAux ()
aux Safety
safety VName
acc [SubExp]
is [SubExp]
vs ADM ()
m
    --
    UserParam {} ->
      ADM (VName, VName) -> ADM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ADM (VName, VName) -> ADM ()) -> ADM (VName, VName) -> ADM ()
forall a b. (a -> b) -> a -> b
$ Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)
commonBasicOp Pat Type
pat StmAux ()
aux BasicOp
e ADM ()
m

vjpOps :: VjpOps
vjpOps :: VjpOps
vjpOps =
  VjpOps
    { vjpLambda :: [Adj] -> [VName] -> Lambda SOACS -> ADM (Lambda SOACS)
vjpLambda = [Adj] -> [VName] -> Lambda SOACS -> ADM (Lambda SOACS)
diffLambda,
      vjpStm :: Stm SOACS -> ADM () -> ADM ()
vjpStm = Stm SOACS -> ADM () -> ADM ()
diffStm,
      vjpBody :: [Adj] -> [VName] -> Body SOACS -> ADM (Body SOACS)
vjpBody = [Adj] -> [VName] -> Body SOACS -> ADM (Body SOACS)
diffBody
    }

diffStm :: Stm SOACS -> ADM () -> ADM ()
diffStm :: Stm SOACS -> ADM () -> ADM ()
diffStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (BasicOp BasicOp
e)) ADM ()
m =
  Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM ()
diffBasicOp Pat Type
Pat (LetDec SOACS)
pat StmAux ()
StmAux (ExpDec SOACS)
aux BasicOp
e ADM ()
m
diffStm stm :: Stm SOACS
stm@(Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
_ (Apply Name
f [(SubExp, Diet)]
args [(RetType SOACS, RetAls)]
_ Safety
_)) ADM ()
m
  | 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
      ADM ()
m

      pat_adj <- VName -> ADM VName
lookupAdjVal (VName -> ADM VName) -> ADM VName -> ADM VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Pat Type -> ADM VName
patName 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)
          convert PrimType
ft PrimType
tt
            | PrimType
ft PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType
tt = PrimExp VName -> PrimExp VName
forall a. a -> a
id
          convert (IntType IntType
ft) (IntType IntType
tt) = ConvOp -> PrimExp VName -> PrimExp VName
forall v. ConvOp -> PrimExp v -> PrimExp v
ConvOpExp (IntType -> IntType -> ConvOp
SExt IntType
ft IntType
tt)
          convert (FloatType FloatType
ft) (FloatType FloatType
tt) = ConvOp -> PrimExp VName -> PrimExp VName
forall v. ConvOp -> PrimExp v -> PrimExp v
ConvOpExp (FloatType -> FloatType -> ConvOp
FPConv FloatType
ft FloatType
tt)
          convert PrimType
Bool (FloatType FloatType
tt) = ConvOp -> PrimExp VName -> PrimExp VName
forall v. ConvOp -> PrimExp v -> PrimExp v
ConvOpExp (FloatType -> ConvOp
BToF FloatType
tt)
          convert (FloatType FloatType
ft) PrimType
Bool = ConvOp -> PrimExp VName -> PrimExp VName
forall v. ConvOp -> PrimExp v -> PrimExp v
ConvOpExp (FloatType -> ConvOp
FToB FloatType
ft)
          convert PrimType
ft PrimType
tt = [Char] -> PrimExp VName -> PrimExp VName
forall a. HasCallStack => [Char] -> a
error ([Char] -> PrimExp VName -> PrimExp VName)
-> [Char] -> PrimExp VName -> PrimExp VName
forall a b. (a -> b) -> a -> b
$ [Char]
"diffStm.convert: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Name, PrimType, PrimType) -> [Char]
forall a. Pretty a => a -> [Char]
prettyString (Name
f, PrimType
ft, PrimType
tt)

      adj_shape <- askShape

      contribs <-
        case pdBuiltin f arg_pes of
          Maybe [PrimExp VName]
Nothing ->
            [Char] -> ADM [VName]
forall a. HasCallStack => [Char] -> a
error ([Char] -> ADM [VName]) -> [Char] -> ADM [VName]
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 ->
            [(PrimExp VName, PrimType)]
-> ((PrimExp VName, PrimType) -> ADM VName) -> ADM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([PrimExp VName] -> [PrimType] -> [(PrimExp VName, PrimType)]
forall a b. [a] -> [b] -> [(a, b)]
zip [PrimExp VName]
derivs [PrimType]
argts) (((PrimExp VName, PrimType) -> ADM VName) -> ADM [VName])
-> ((PrimExp VName, PrimType) -> ADM VName) -> ADM [VName]
forall a b. (a -> b) -> a -> b
$ \(PrimExp VName
deriv, PrimType
argt) ->
              Name -> Exp (Rep ADM) -> ADM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"apply_contrib" (Exp SOACS -> ADM VName)
-> ((Solo SubExp -> ADM (Exp SOACS)) -> ADM (Exp SOACS))
-> (Solo SubExp -> ADM (Exp SOACS))
-> ADM VName
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Shape
-> Solo SubExp
-> (Solo SubExp -> ADM (Exp SOACS))
-> ADM (Exp SOACS)
forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Rep m ~ SOACS, Traversable f) =>
Shape -> f SubExp -> (f SubExp -> m (Exp SOACS)) -> m (Exp SOACS)
mapNest Shape
adj_shape (SubExp -> Solo SubExp
forall a. a -> Solo a
MkSolo (VName -> SubExp
Var VName
pat_adj)) ((Solo SubExp -> ADM (Exp SOACS)) -> ADM VName)
-> (Solo SubExp -> ADM (Exp SOACS)) -> ADM VName
forall a b. (a -> b) -> a -> b
$
                \(MkSolo SubExp
pat_adj') ->
                  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 -> PrimType -> PrimExp VName -> PrimExp VName
convert PrimType
ret PrimType
argt (PrimExp VName -> PrimExp VName) -> PrimExp VName -> PrimExp VName
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
ret SubExp
pat_adj' PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. PrimExp v -> PrimExp v -> PrimExp v
~*~ PrimExp VName
deriv

      zipWithM_ updateSubExpAdj (map fst args) contribs
diffStm stm :: Stm SOACS
stm@(Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
_ (Match [SubExp]
ses [Case (Body SOACS)]
cases Body SOACS
defbody MatchDec (BranchType SOACS)
_)) ADM ()
m = do
  Stm (Rep ADM) -> ADM ()
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm Stm (Rep ADM)
Stm SOACS
stm
  ADM ()
m
  ADM () -> ADM ()
forall a. ADM a -> ADM a
returnSweepCode (ADM () -> ADM ()) -> ADM () -> ADM ()
forall a b. (a -> b) -> a -> b
$ do
    let cases_free :: [Names]
cases_free = (Case (Body SOACS) -> Names) -> [Case (Body SOACS)] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map Case (Body SOACS) -> Names
forall a. FreeIn a => a -> Names
freeIn [Case (Body SOACS)]
cases
        defbody_free :: Names
defbody_free = Body SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Body SOACS
defbody
        branches_free :: [VName]
branches_free = Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat ([Names] -> Names) -> [Names] -> Names
forall a b. (a -> b) -> a -> b
$ Names
defbody_free Names -> [Names] -> [Names]
forall a. a -> [a] -> [a]
: [Names]
cases_free

    adjs <- (VName -> ADM Adj) -> [VName] -> ADM [Adj]
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 Adj
lookupAdj ([VName] -> ADM [Adj]) -> [VName] -> ADM [Adj]
forall a b. (a -> b) -> a -> b
$ Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
Pat (LetDec SOACS)
pat

    branches_free_adj <-
      ( pure . takeLast (length branches_free)
          <=< letTupExp "branch_adj"
          <=< renameExp
      )
        =<< eMatch
          ses
          (map (fmap $ diffBody adjs branches_free) cases)
          (diffBody adjs branches_free defbody)
    -- See Note [Array Adjoints of Match]
    forM_ (zip branches_free branches_free_adj) $ \(VName
v, VName
v_adj) ->
      VName -> VName -> ADM ()
insAdj VName
v (VName -> ADM ()) -> ADM VName -> ADM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> ADM VName
copyIfArray VName
v_adj
diffStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Op Op SOACS
soac)) ADM ()
m =
  -- We add the attributes from 'aux' to every SOAC (but only SOAC) produced. We
  -- could do this on *every* stm, but it would be very verbose.
  (Stms (Rep ADM) -> Stms (Rep ADM)) -> ADM () -> ADM ()
forall (m :: * -> *) a.
MonadBuilder m =>
(Stms (Rep m) -> Stms (Rep m)) -> m a -> m a
censorStms ((Stm SOACS -> Stm SOACS) -> Stms SOACS -> Stms SOACS
forall a b. (a -> b) -> Seq a -> Seq b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Stm SOACS -> Stm SOACS
addAttrs) (ADM () -> ADM ()) -> ADM () -> ADM ()
forall a b. (a -> b) -> a -> b
$ VjpOps -> Pat Type -> StmAux () -> SOAC SOACS -> ADM () -> ADM ()
vjpSOAC VjpOps
vjpOps Pat Type
Pat (LetDec SOACS)
pat StmAux ()
StmAux (ExpDec SOACS)
aux Op SOACS
SOAC SOACS
soac ADM ()
m
  where
    addAttrs :: Stm SOACS -> Stm SOACS
addAttrs Stm SOACS
stm
      | Op Op SOACS
_ <- Stm SOACS -> Exp SOACS
forall rep. Stm rep -> Exp rep
stmExp Stm SOACS
stm =
          Attrs -> Stm SOACS -> Stm SOACS
forall rep. Attrs -> Stm rep -> Stm rep
attribute (StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux ()
StmAux (ExpDec SOACS)
aux) Stm SOACS
stm
      | Bool
otherwise = Stm SOACS
stm
diffStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux loop :: Exp SOACS
loop@Loop {}) ADM ()
m =
  (Stms SOACS -> ADM ())
-> Pat Type -> StmAux () -> Exp SOACS -> ADM () -> ADM ()
diffLoop Stms SOACS -> ADM ()
diffStms Pat Type
Pat (LetDec SOACS)
pat StmAux ()
StmAux (ExpDec SOACS)
aux Exp SOACS
loop ADM ()
m
-- See Note [Adjoints of accumulators]
diffStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (WithAcc [WithAccInput SOACS]
inputs Lambda SOACS
lam)) ADM ()
m =
  VjpOps
-> Pat Type
-> StmAux ()
-> [WithAccInput SOACS]
-> Lambda SOACS
-> ADM ()
-> ADM ()
diffWithAcc VjpOps
vjpOps Pat Type
Pat (LetDec SOACS)
pat StmAux ()
StmAux (ExpDec SOACS)
aux [WithAccInput SOACS]
inputs Lambda SOACS
lam ADM ()
m
diffStm Stm SOACS
stm ADM ()
_ = [Char] -> ADM ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> ADM ()) -> [Char] -> ADM ()
forall a b. (a -> b) -> a -> b
$ [Char]
"diffStm unhandled:\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Stm SOACS -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Stm SOACS
stm

diffStms :: Stms SOACS -> ADM ()
diffStms :: Stms SOACS -> ADM ()
diffStms Stms SOACS
all_stms
  | Just (Stm SOACS
stm, Stms SOACS
stms) <- Stms SOACS -> Maybe (Stm SOACS, Stms SOACS)
forall rep. Stms rep -> Maybe (Stm rep, Stms rep)
stmsHead Stms SOACS
all_stms = do
      (subst, copy_stms) <- Stm SOACS -> ADM (Substitutions, Stms SOACS)
copyConsumedArrsInStm Stm SOACS
stm
      let (stm', stms') = substituteNames subst (stm, stms)
      diffStms copy_stms >> diffStm stm' (diffStms stms')
      forM_ (M.toList subst) $ \(VName
from, VName
to) ->
        VName -> Adj -> ADM ()
setAdj VName
from (Adj -> ADM ()) -> ADM Adj -> ADM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> ADM Adj
lookupAdj VName
to
  | Bool
otherwise =
      () -> ADM ()
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Preprocess statements before differentiating.
-- For now, it's just stripmining.
preprocess :: Stms SOACS -> ADM (Stms SOACS)
preprocess :: Stms SOACS -> ADM (Stms SOACS)
preprocess = Stms SOACS -> ADM (Stms SOACS)
stripmineStms

diffBody :: [Adj] -> [VName] -> Body SOACS -> ADM (Body SOACS)
diffBody :: [Adj] -> [VName] -> Body SOACS -> ADM (Body SOACS)
diffBody [Adj]
res_adjs [VName]
get_adjs_for (Body () Stms SOACS
stms [SubExpRes]
res) = ADM (Body SOACS) -> ADM (Body SOACS)
forall a. ADM a -> ADM a
subAD (ADM (Body SOACS) -> ADM (Body SOACS))
-> ADM (Body SOACS) -> ADM (Body SOACS)
forall a b. (a -> b) -> a -> b
$
  ADM (Body SOACS) -> ADM (Body SOACS)
forall a. ADM a -> ADM a
subSubsts (ADM (Body SOACS) -> ADM (Body SOACS))
-> ADM (Body SOACS) -> ADM (Body SOACS)
forall a b. (a -> b) -> a -> b
$ do
    let onResult :: SubExpRes -> Adj -> ADM ()
onResult (SubExpRes Certs
_ (Constant PrimValue
_)) Adj
_ = () -> ADM ()
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        onResult (SubExpRes Certs
_ (Var VName
v)) Adj
v_adj = ADM () -> ADM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ADM () -> ADM ()) -> ADM () -> ADM ()
forall a b. (a -> b) -> a -> b
$ VName -> VName -> ADM ()
updateAdj VName
v (VName -> ADM ()) -> ADM VName -> ADM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Adj -> ADM VName
adjVal Adj
v_adj
    (adjs, stms') <- ADM [VName] -> ADM ([VName], Stms (Rep ADM))
forall a. ADM a -> ADM (a, Stms (Rep ADM))
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms (ADM [VName] -> ADM ([VName], Stms (Rep ADM)))
-> ADM [VName] -> ADM ([VName], Stms (Rep ADM))
forall a b. (a -> b) -> a -> b
$ do
      (SubExpRes -> Adj -> ADM ()) -> [SubExpRes] -> [Adj] -> ADM ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ SubExpRes -> Adj -> ADM ()
onResult (Int -> [SubExpRes] -> [SubExpRes]
forall a. Int -> [a] -> [a]
takeLast ([Adj] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Adj]
res_adjs) [SubExpRes]
res) [Adj]
res_adjs
      Stms SOACS -> ADM ()
diffStms (Stms SOACS -> ADM ()) -> ADM (Stms SOACS) -> ADM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Stms SOACS -> ADM (Stms SOACS)
preprocess Stms SOACS
stms
      (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
lookupAdjVal [VName]
get_adjs_for
    pure $ Body () stms' $ res <> varsRes adjs

diffLambda :: [Adj] -> [VName] -> Lambda SOACS -> ADM (Lambda SOACS)
diffLambda :: [Adj] -> [VName] -> Lambda SOACS -> ADM (Lambda SOACS)
diffLambda [Adj]
res_adjs [VName]
get_adjs_for (Lambda [LParam SOACS]
params [Type]
_ Body SOACS
body) =
  [LParam (Rep ADM)] -> ADM [SubExpRes] -> ADM (Lambda (Rep ADM))
forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m [SubExpRes] -> m (Lambda (Rep m))
mkLambda [LParam (Rep ADM)]
[LParam SOACS]
params (ADM [SubExpRes] -> ADM (Lambda (Rep ADM)))
-> ADM [SubExpRes] -> ADM (Lambda (Rep ADM))
forall a b. (a -> b) -> a -> b
$ do
    res <- Body (Rep ADM) -> ADM [SubExpRes]
Body SOACS -> ADM [SubExpRes]
forall (m :: * -> *).
MonadBuilder m =>
Body (Rep m) -> m [SubExpRes]
bodyBind (Body SOACS -> ADM [SubExpRes])
-> ADM (Body SOACS) -> ADM [SubExpRes]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Adj] -> [VName] -> Body SOACS -> ADM (Body SOACS)
diffBody [Adj]
res_adjs [VName]
get_adjs_for Body SOACS
body
    pure $ takeLast (length get_adjs_for) res

revVJP ::
  (MonadFreshNames m) =>
  Scope SOACS ->
  Shape ->
  Attrs ->
  Lambda SOACS ->
  m (Lambda SOACS)
revVJP :: forall (m :: * -> *).
MonadFreshNames m =>
Scope SOACS -> Shape -> Attrs -> Lambda SOACS -> m (Lambda SOACS)
revVJP Scope SOACS
scope Shape
shape Attrs
attrs (Lambda [LParam SOACS]
params [Type]
ts Body SOACS
body) = do
  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 Scope SOACS -> Scope SOACS -> Scope SOACS
forall a. Semigroup a => a -> a -> a
<> [Param Type] -> Scope SOACS
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [Param Type]
[LParam SOACS]
params) (ADM (Lambda SOACS) -> m (Lambda SOACS))
-> ADM (Lambda SOACS) -> m (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ do
    adj_shape <- ADM Shape
askShape
    params_adj <- forM (zip (map resSubExp (bodyResult body)) ts) $ \(SubExp
se, Type
t) ->
      Attrs -> VName -> Type -> Param Type
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
forall a. Monoid a => a
mempty
        (VName -> Type -> Param Type)
-> ADM VName -> ADM (Type -> Param Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ADM VName -> (VName -> ADM VName) -> Maybe VName -> ADM VName
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Name -> ADM VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"const_res_adj") VName -> ADM VName
adjVName (SubExp -> Maybe VName
subExpVar SubExp
se)
        ADM (Type -> Param Type) -> ADM Type -> ADM (Param Type)
forall a b. ADM (a -> b) -> ADM a -> ADM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> ADM Type
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type
t Type -> Shape -> Type
`arrayOfShape` Shape
adj_shape)

    body' <-
      localScope (scopeOfLParams params_adj) $
        diffBody
          (map adjFromParam params_adj)
          (map paramName params)
          body

    pure $ Lambda (params ++ params_adj) (ts <> map paramType params) body'