{-# 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.Monad
  ( ADM,
    RState (..),
    REnv,
    runADM,
    Adj (..),
    InBounds (..),
    Sparse (..),
    askShape,
    adjFromParam,
    adjFromVar,
    lookupAdj,
    lookupAdjVal,
    adjVal,
    updateAdj,
    updateSubExpAdj,
    updateAdjSlice,
    updateAdjIndex,
    setAdj,
    insAdj,
    adjsReps,
    --
    copyConsumedArrsInStm,
    copyConsumedArrsInBody,
    addSubstitution,
    returnSweepCode,
    --
    adjVName,
    subAD,
    noAdjsFor,
    subSubsts,
    isActive,
    --
    tabNest,
    oneExp,
    zeroExp,
    zeroArray,
    unitAdjOfType,
    addLambda,
    vecOpExp,
    --
    VjpOps (..),
    --
    setLoopTape,
    lookupLoopTape,
    substLoopTape,
    renameLoopTape,
    --
    locallyNonvector,
    vecToInner,
    vecSlice,
    withOperatorAccs,
    isOperatorAcc,
  )
where

import Control.Monad
import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.Bifunctor (second)
import Data.Map qualified as M
import Data.Maybe
import Futhark.AD.Shared
import Futhark.Analysis.Alias qualified as Alias
import Futhark.Analysis.PrimExp.Convert
import Futhark.Builder
import Futhark.IR.Aliases (consumedInStms)
import Futhark.IR.Prop.Aliases
import Futhark.IR.SOACS
import Futhark.Tools
import Futhark.Transform.Substitute
import Futhark.Util (chunks)

zeroExp :: Type -> Exp rep
zeroExp :: forall rep. TypeBase Shape NoMode -> Exp rep
zeroExp (Prim PrimType
pt) =
  BasicOp -> Exp rep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
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 rep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
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 TypeBase Shape NoMode
t = [Char] -> Exp rep
forall a. HasCallStack => [Char] -> a
error ([Char] -> Exp rep) -> [Char] -> Exp rep
forall a b. (a -> b) -> a -> b
$ [Char]
"zeroExp: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ TypeBase Shape NoMode -> [Char]
forall a. Pretty a => a -> [Char]
prettyString TypeBase Shape NoMode
t

onePrim :: PrimType -> PrimValue
onePrim :: PrimType -> PrimValue
onePrim (IntType IntType
it) = IntValue -> PrimValue
IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Int -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it (Int
1 :: Int)
onePrim (FloatType FloatType
ft) = FloatValue -> PrimValue
FloatValue (FloatValue -> PrimValue) -> FloatValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ FloatType -> Double -> FloatValue
forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft (Double
1 :: Double)
onePrim PrimType
Bool = Bool -> PrimValue
BoolValue Bool
True
onePrim PrimType
Unit = PrimValue
UnitValue

oneExp :: Type -> Exp rep
oneExp :: forall rep. TypeBase Shape NoMode -> Exp rep
oneExp (Prim PrimType
t) = BasicOp -> Exp rep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimValue
onePrim PrimType
t
oneExp (Array PrimType
pt Shape
shape NoMode
_) =
  BasicOp -> Exp rep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
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
onePrim PrimType
pt
oneExp TypeBase Shape NoMode
t = [Char] -> Exp rep
forall a. HasCallStack => [Char] -> a
error ([Char] -> Exp rep) -> [Char] -> Exp rep
forall a b. (a -> b) -> a -> b
$ [Char]
"oneExp: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ TypeBase Shape NoMode -> [Char]
forall a. Pretty a => a -> [Char]
prettyString TypeBase Shape NoMode
t

-- | Whether 'Sparse' should check bounds or assume they are correct.
-- The latter results in simpler code.
data InBounds
  = -- | If a SubExp is provided, it references a boolean that is true
    -- when in-bounds.
    CheckBounds (Maybe SubExp)
  | -- | Assume that these are always in-bounds.
    AssumeBounds
  | -- | Dynamically these will always fail, so don't bother
    -- generating code for the update.  This is only needed to ensure
    -- a consistent representation of sparse Jacobians.
    OutOfBounds
  deriving (InBounds -> InBounds -> Bool
(InBounds -> InBounds -> Bool)
-> (InBounds -> InBounds -> Bool) -> Eq InBounds
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: InBounds -> InBounds -> Bool
== :: InBounds -> InBounds -> Bool
$c/= :: InBounds -> InBounds -> Bool
/= :: InBounds -> InBounds -> Bool
Eq, Eq InBounds
Eq InBounds =>
(InBounds -> InBounds -> Ordering)
-> (InBounds -> InBounds -> Bool)
-> (InBounds -> InBounds -> Bool)
-> (InBounds -> InBounds -> Bool)
-> (InBounds -> InBounds -> Bool)
-> (InBounds -> InBounds -> InBounds)
-> (InBounds -> InBounds -> InBounds)
-> Ord InBounds
InBounds -> InBounds -> Bool
InBounds -> InBounds -> Ordering
InBounds -> InBounds -> InBounds
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: InBounds -> InBounds -> Ordering
compare :: InBounds -> InBounds -> Ordering
$c< :: InBounds -> InBounds -> Bool
< :: InBounds -> InBounds -> Bool
$c<= :: InBounds -> InBounds -> Bool
<= :: InBounds -> InBounds -> Bool
$c> :: InBounds -> InBounds -> Bool
> :: InBounds -> InBounds -> Bool
$c>= :: InBounds -> InBounds -> Bool
>= :: InBounds -> InBounds -> Bool
$cmax :: InBounds -> InBounds -> InBounds
max :: InBounds -> InBounds -> InBounds
$cmin :: InBounds -> InBounds -> InBounds
min :: InBounds -> InBounds -> InBounds
Ord, Int -> InBounds -> [Char] -> [Char]
[InBounds] -> [Char] -> [Char]
InBounds -> [Char]
(Int -> InBounds -> [Char] -> [Char])
-> (InBounds -> [Char])
-> ([InBounds] -> [Char] -> [Char])
-> Show InBounds
forall a.
(Int -> a -> [Char] -> [Char])
-> (a -> [Char]) -> ([a] -> [Char] -> [Char]) -> Show a
$cshowsPrec :: Int -> InBounds -> [Char] -> [Char]
showsPrec :: Int -> InBounds -> [Char] -> [Char]
$cshow :: InBounds -> [Char]
show :: InBounds -> [Char]
$cshowList :: [InBounds] -> [Char] -> [Char]
showList :: [InBounds] -> [Char] -> [Char]
Show)

-- | A symbolic representation of an array that is all zeroes, except
-- at certain indexes.
data Sparse = Sparse
  { -- | The full shape of the array (including any vector dimensions, which are
    -- stored in sparseVecDims).
    Sparse -> Shape
sparseShape :: Shape,
    -- | Element type of the array.
    Sparse -> PrimType
sparseType :: PrimType,
    -- | Number of leading dimensions that are \"vector\" dimensions, due to
    -- vector AD. These are not indexed by the sparse index, but are present in
    -- the values. When zero, this is the ordinary non-vector case. This is
    -- equivalent to the rank of `askShape`, but it is convenient to store it
    -- here as well.
    Sparse -> Int
sparseVecDims :: Int,
    -- | Locations and values of nonzero values.  Indexes may be
    -- negative, in which case the value is ignored (unless
    -- 'AssumeBounds' is used).
    Sparse -> [(InBounds, SubExp, SubExp)]
sparseIdxVals :: [(InBounds, SubExp, SubExp)]
  }
  deriving (Sparse -> Sparse -> Bool
(Sparse -> Sparse -> Bool)
-> (Sparse -> Sparse -> Bool) -> Eq Sparse
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Sparse -> Sparse -> Bool
== :: Sparse -> Sparse -> Bool
$c/= :: Sparse -> Sparse -> Bool
/= :: Sparse -> Sparse -> Bool
Eq, Eq Sparse
Eq Sparse =>
(Sparse -> Sparse -> Ordering)
-> (Sparse -> Sparse -> Bool)
-> (Sparse -> Sparse -> Bool)
-> (Sparse -> Sparse -> Bool)
-> (Sparse -> Sparse -> Bool)
-> (Sparse -> Sparse -> Sparse)
-> (Sparse -> Sparse -> Sparse)
-> Ord Sparse
Sparse -> Sparse -> Bool
Sparse -> Sparse -> Ordering
Sparse -> Sparse -> Sparse
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Sparse -> Sparse -> Ordering
compare :: Sparse -> Sparse -> Ordering
$c< :: Sparse -> Sparse -> Bool
< :: Sparse -> Sparse -> Bool
$c<= :: Sparse -> Sparse -> Bool
<= :: Sparse -> Sparse -> Bool
$c> :: Sparse -> Sparse -> Bool
> :: Sparse -> Sparse -> Bool
$c>= :: Sparse -> Sparse -> Bool
>= :: Sparse -> Sparse -> Bool
$cmax :: Sparse -> Sparse -> Sparse
max :: Sparse -> Sparse -> Sparse
$cmin :: Sparse -> Sparse -> Sparse
min :: Sparse -> Sparse -> Sparse
Ord, Int -> Sparse -> [Char] -> [Char]
[Sparse] -> [Char] -> [Char]
Sparse -> [Char]
(Int -> Sparse -> [Char] -> [Char])
-> (Sparse -> [Char])
-> ([Sparse] -> [Char] -> [Char])
-> Show Sparse
forall a.
(Int -> a -> [Char] -> [Char])
-> (a -> [Char]) -> ([a] -> [Char] -> [Char]) -> Show a
$cshowsPrec :: Int -> Sparse -> [Char] -> [Char]
showsPrec :: Int -> Sparse -> [Char] -> [Char]
$cshow :: Sparse -> [Char]
show :: Sparse -> [Char]
$cshowList :: [Sparse] -> [Char] -> [Char]
showList :: [Sparse] -> [Char] -> [Char]
Show)

-- | The adjoint of a variable.
data Adj
  = AdjSparse Sparse
  | AdjVal SubExp
  | AdjZero Shape PrimType
  deriving (Adj -> Adj -> Bool
(Adj -> Adj -> Bool) -> (Adj -> Adj -> Bool) -> Eq Adj
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Adj -> Adj -> Bool
== :: Adj -> Adj -> Bool
$c/= :: Adj -> Adj -> Bool
/= :: Adj -> Adj -> Bool
Eq, Eq Adj
Eq Adj =>
(Adj -> Adj -> Ordering)
-> (Adj -> Adj -> Bool)
-> (Adj -> Adj -> Bool)
-> (Adj -> Adj -> Bool)
-> (Adj -> Adj -> Bool)
-> (Adj -> Adj -> Adj)
-> (Adj -> Adj -> Adj)
-> Ord Adj
Adj -> Adj -> Bool
Adj -> Adj -> Ordering
Adj -> Adj -> Adj
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Adj -> Adj -> Ordering
compare :: Adj -> Adj -> Ordering
$c< :: Adj -> Adj -> Bool
< :: Adj -> Adj -> Bool
$c<= :: Adj -> Adj -> Bool
<= :: Adj -> Adj -> Bool
$c> :: Adj -> Adj -> Bool
> :: Adj -> Adj -> Bool
$c>= :: Adj -> Adj -> Bool
>= :: Adj -> Adj -> Bool
$cmax :: Adj -> Adj -> Adj
max :: Adj -> Adj -> Adj
$cmin :: Adj -> Adj -> Adj
min :: Adj -> Adj -> Adj
Ord, Int -> Adj -> [Char] -> [Char]
[Adj] -> [Char] -> [Char]
Adj -> [Char]
(Int -> Adj -> [Char] -> [Char])
-> (Adj -> [Char]) -> ([Adj] -> [Char] -> [Char]) -> Show Adj
forall a.
(Int -> a -> [Char] -> [Char])
-> (a -> [Char]) -> ([a] -> [Char] -> [Char]) -> Show a
$cshowsPrec :: Int -> Adj -> [Char] -> [Char]
showsPrec :: Int -> Adj -> [Char] -> [Char]
$cshow :: Adj -> [Char]
show :: Adj -> [Char]
$cshowList :: [Adj] -> [Char] -> [Char]
showList :: [Adj] -> [Char] -> [Char]
Show)

instance Substitute Adj where
  substituteNames :: Substitutions -> Adj -> Adj
substituteNames Substitutions
m (AdjVal (Var VName
v)) = SubExp -> Adj
AdjVal (SubExp -> Adj) -> SubExp -> Adj
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Substitutions -> VName -> VName
forall a. Substitute a => Substitutions -> a -> a
substituteNames Substitutions
m VName
v
  substituteNames Substitutions
_ Adj
adj = Adj
adj

-- | Create an array of the given shape and element type consisting of zeroes.
-- The shape may be empty, meaning this function can (despite its name) also
-- create non-arrays.
zeroArray :: (MonadBuilder m) => Shape -> Type -> m VName
zeroArray :: forall (m :: * -> *).
MonadBuilder m =>
Shape -> TypeBase Shape NoMode -> m VName
zeroArray Shape
shape TypeBase Shape NoMode
t
  | Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
shape Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 =
      Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"zero" (Exp (Rep m) -> m VName) -> Exp (Rep m) -> m VName
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoMode -> Exp (Rep m)
forall rep. TypeBase Shape NoMode -> Exp rep
zeroExp TypeBase Shape NoMode
t
  | Bool
otherwise = do
      zero <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"zero" (Exp (Rep m) -> m SubExp) -> Exp (Rep m) -> m SubExp
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoMode -> Exp (Rep m)
forall rep. TypeBase Shape NoMode -> Exp rep
zeroExp TypeBase Shape NoMode
t
      attributing (oneAttr "sequential") $
        letExp "zeroes_" . BasicOp $
          Replicate shape zero

sparseArray :: (MonadBuilder m, Rep m ~ SOACS) => Sparse -> m VName
sparseArray :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ SOACS) =>
Sparse -> m VName
sparseArray (Sparse Shape
shape PrimType
t Int
vec_dims [(InBounds, SubExp, SubExp)]
ivs) = do
  (VName -> [(InBounds, SubExp, SubExp)] -> m VName)
-> [(InBounds, SubExp, SubExp)] -> VName -> m VName
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((VName -> (InBounds, SubExp, SubExp) -> m VName)
-> VName -> [(InBounds, SubExp, SubExp)] -> m VName
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM VName -> (InBounds, SubExp, SubExp) -> m VName
f) [(InBounds, SubExp, SubExp)]
ivs (VName -> m VName) -> m VName -> m VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Shape -> TypeBase Shape NoMode -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Shape -> TypeBase Shape NoMode -> m VName
zeroArray Shape
shape (PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
t)
  where
    arr_t :: TypeBase Shape NoMode
arr_t = PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
t TypeBase Shape NoMode -> Shape -> TypeBase Shape NoMode
`arrayOfShape` Shape
shape
    vec_slice :: [DimIndex SubExp]
vec_slice = (SubExp -> DimIndex SubExp) -> [SubExp] -> [DimIndex SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimIndex SubExp
sliceDim ([SubExp] -> [DimIndex SubExp]) -> [SubExp] -> [DimIndex SubExp]
forall a b. (a -> b) -> a -> b
$ Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take Int
vec_dims ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape
    f :: VName -> (InBounds, SubExp, SubExp) -> m VName
f VName
arr (InBounds
check, SubExp
i, SubExp
se) = do
      let stm :: Safety -> m VName
stm Safety
s =
            Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"sparse" (Exp SOACS -> m VName)
-> (BasicOp -> Exp SOACS) -> BasicOp -> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> m VName) -> BasicOp -> m VName
forall a b. (a -> b) -> a -> b
$
              Safety -> VName -> Slice SubExp -> SubExp -> BasicOp
Update Safety
s VName
arr (TypeBase Shape NoMode -> [DimIndex SubExp] -> Slice SubExp
fullSlice TypeBase Shape NoMode
arr_t ([DimIndex SubExp]
vec_slice [DimIndex SubExp] -> [DimIndex SubExp] -> [DimIndex SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix SubExp
i])) SubExp
se
      case InBounds
check of
        InBounds
AssumeBounds -> Safety -> m VName
stm Safety
Unsafe
        CheckBounds Maybe SubExp
_ -> Safety -> m VName
stm Safety
Safe
        InBounds
OutOfBounds -> VName -> m VName
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
arr

adjFromVar :: VName -> Adj
adjFromVar :: VName -> Adj
adjFromVar = SubExp -> Adj
AdjVal (SubExp -> Adj) -> (VName -> SubExp) -> VName -> Adj
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var

adjFromParam :: Param t -> Adj
adjFromParam :: forall t. Param t -> Adj
adjFromParam = VName -> Adj
adjFromVar (VName -> Adj) -> (Param t -> VName) -> Param t -> Adj
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param t -> VName
forall dec. Param dec -> VName
paramName

unitAdjOfType :: Type -> ADM Adj
unitAdjOfType :: TypeBase Shape NoMode -> ADM Adj
unitAdjOfType TypeBase Shape NoMode
t = SubExp -> Adj
AdjVal (SubExp -> Adj) -> ADM SubExp -> ADM Adj
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Exp (Rep ADM) -> ADM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"adj_unit" (TypeBase Shape NoMode -> Exp SOACS
forall rep. TypeBase Shape NoMode -> Exp rep
oneExp TypeBase Shape NoMode
t)

-- | The values representing an adjoint in symbolic form.  This is
-- used for when we wish to return an Adj from a Body or similar
-- without forcing manifestation.  Also returns a function for
-- reassembling the Adj from a new representation (the list must have
-- the same length).
adjRep :: Adj -> ([SubExp], [SubExp] -> Adj)
adjRep :: Adj -> ([SubExp], [SubExp] -> Adj)
adjRep (AdjVal SubExp
se) = ([SubExp
se], \[SubExp
se'] -> SubExp -> Adj
AdjVal SubExp
se')
adjRep (AdjZero Shape
shape PrimType
pt) = ([], \[] -> Shape -> PrimType -> Adj
AdjZero Shape
shape PrimType
pt)
adjRep (AdjSparse (Sparse Shape
shape PrimType
pt Int
vd [(InBounds, SubExp, SubExp)]
ivs)) =
  (((InBounds, SubExp, SubExp) -> [SubExp])
-> [(InBounds, SubExp, SubExp)] -> [SubExp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (InBounds, SubExp, SubExp) -> [SubExp]
forall {a} {a}. (a, a, a) -> [a]
ivRep [(InBounds, SubExp, SubExp)]
ivs, Sparse -> Adj
AdjSparse (Sparse -> Adj) -> ([SubExp] -> Sparse) -> [SubExp] -> Adj
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape -> PrimType -> Int -> [(InBounds, SubExp, SubExp)] -> Sparse
Sparse Shape
shape PrimType
pt Int
vd ([(InBounds, SubExp, SubExp)] -> Sparse)
-> ([SubExp] -> [(InBounds, SubExp, SubExp)]) -> [SubExp] -> Sparse
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(InBounds, SubExp, SubExp)]
-> [SubExp] -> [(InBounds, SubExp, SubExp)]
forall {b} {c} {c}. [(InBounds, b, c)] -> [c] -> [(InBounds, c, c)]
repIvs [(InBounds, SubExp, SubExp)]
ivs)
  where
    ivRep :: (a, a, a) -> [a]
ivRep (a
_, a
i, a
v) = [a
i, a
v]
    repIvs :: [(InBounds, b, c)] -> [c] -> [(InBounds, c, c)]
repIvs ((InBounds
check, b
_, c
_) : [(InBounds, b, c)]
ivs') (c
i : c
v : [c]
ses) =
      (InBounds
check', c
i, c
v) (InBounds, c, c) -> [(InBounds, c, c)] -> [(InBounds, c, c)]
forall a. a -> [a] -> [a]
: [(InBounds, b, c)] -> [c] -> [(InBounds, c, c)]
repIvs [(InBounds, b, c)]
ivs' [c]
ses
      where
        check' :: InBounds
check' = case InBounds
check of
          InBounds
AssumeBounds -> InBounds
AssumeBounds
          CheckBounds Maybe SubExp
b -> Maybe SubExp -> InBounds
CheckBounds Maybe SubExp
b
          InBounds
OutOfBounds -> Maybe SubExp -> InBounds
CheckBounds (SubExp -> Maybe SubExp
forall a. a -> Maybe a
Just (Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False)) -- sic!
    repIvs [(InBounds, b, c)]
_ [c]
_ = []

-- | Conveniently convert a list of Adjs to their representation, as
-- well as produce a function for converting back.
adjsReps :: [Adj] -> ([SubExp], [SubExp] -> [Adj])
adjsReps :: [Adj] -> ([SubExp], [SubExp] -> [Adj])
adjsReps [Adj]
adjs =
  let ([[SubExp]]
reps, [[SubExp] -> Adj]
fs) = [([SubExp], [SubExp] -> Adj)] -> ([[SubExp]], [[SubExp] -> Adj])
forall a b. [(a, b)] -> ([a], [b])
unzip ([([SubExp], [SubExp] -> Adj)] -> ([[SubExp]], [[SubExp] -> Adj]))
-> [([SubExp], [SubExp] -> Adj)] -> ([[SubExp]], [[SubExp] -> Adj])
forall a b. (a -> b) -> a -> b
$ (Adj -> ([SubExp], [SubExp] -> Adj))
-> [Adj] -> [([SubExp], [SubExp] -> Adj)]
forall a b. (a -> b) -> [a] -> [b]
map Adj -> ([SubExp], [SubExp] -> Adj)
adjRep [Adj]
adjs
   in ([[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[SubExp]]
reps, (([SubExp] -> Adj) -> [SubExp] -> Adj)
-> [[SubExp] -> Adj] -> [[SubExp]] -> [Adj]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ([SubExp] -> Adj) -> [SubExp] -> Adj
forall a b. (a -> b) -> a -> b
($) [[SubExp] -> Adj]
fs ([[SubExp]] -> [Adj])
-> ([SubExp] -> [[SubExp]]) -> [SubExp] -> [Adj]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [SubExp] -> [[SubExp]]
forall a. [Int] -> [a] -> [[a]]
chunks (([SubExp] -> Int) -> [[SubExp]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[SubExp]]
reps))

data RState = RState
  { RState -> Map VName Adj
stateAdjs :: M.Map VName Adj,
    RState -> Substitutions
stateLoopTape :: Substitutions,
    RState -> Substitutions
stateSubsts :: Substitutions,
    RState -> VNameSource
stateNameSource :: VNameSource
  }

data REnv = REnv
  { REnv -> Shape
envAdjShape :: Shape,
    REnv -> Attrs
envAttrs :: Attrs,
    -- | The certificates of those accumulators that have a combining operator,
    -- and so do not behave like overwrites.  See Note [Adjoints of
    -- accumulators].
    REnv -> Names
envOperatorAccs :: Names
  }

newtype ADM a = ADM (BuilderT SOACS (ReaderT REnv (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 REnv,
      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 REnv (State RState)) (ExpDec (Rep ADM))
-> ADM (ExpDec (Rep ADM))
forall a. BuilderT SOACS (ReaderT REnv (State RState)) a -> ADM a
ADM (BuilderT SOACS (ReaderT REnv (State RState)) (ExpDec (Rep ADM))
 -> ADM (ExpDec (Rep ADM)))
-> BuilderT SOACS (ReaderT REnv (State RState)) (ExpDec (Rep ADM))
-> ADM (ExpDec (Rep ADM))
forall a b. (a -> b) -> a -> b
$ Pat (LetDec (Rep (BuilderT SOACS (ReaderT REnv (State RState)))))
-> Exp (Rep (BuilderT SOACS (ReaderT REnv (State RState))))
-> BuilderT
     SOACS
     (ReaderT REnv (State RState))
     (ExpDec (Rep (BuilderT SOACS (ReaderT REnv (State RState)))))
forall (m :: * -> *).
MonadBuilder m =>
Pat (LetDec (Rep m)) -> Exp (Rep m) -> m (ExpDec (Rep m))
mkExpDecM Pat (LetDec (Rep (BuilderT SOACS (ReaderT REnv (State RState)))))
Pat (LetDec (Rep ADM))
pat Exp (Rep (BuilderT SOACS (ReaderT REnv (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 REnv (State RState)) (GBody (Rep ADM) res)
-> ADM (GBody (Rep ADM) res)
forall a. BuilderT SOACS (ReaderT REnv (State RState)) a -> ADM a
ADM (BuilderT SOACS (ReaderT REnv (State RState)) (GBody (Rep ADM) res)
 -> ADM (GBody (Rep ADM) res))
-> BuilderT
     SOACS (ReaderT REnv (State RState)) (GBody (Rep ADM) res)
-> ADM (GBody (Rep ADM) res)
forall a b. (a -> b) -> a -> b
$ Stms (Rep (BuilderT SOACS (ReaderT REnv (State RState))))
-> [res]
-> BuilderT
     SOACS
     (ReaderT REnv (State RState))
     (GBody (Rep (BuilderT SOACS (ReaderT REnv (State RState)))) res)
forall res.
IsResult res =>
Stms (Rep (BuilderT SOACS (ReaderT REnv (State RState))))
-> [res]
-> BuilderT
     SOACS
     (ReaderT REnv (State RState))
     (GBody (Rep (BuilderT SOACS (ReaderT REnv (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 REnv (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 REnv (State RState)) (Stm (Rep ADM))
-> ADM (Stm (Rep ADM))
forall a. BuilderT SOACS (ReaderT REnv (State RState)) a -> ADM a
ADM (BuilderT SOACS (ReaderT REnv (State RState)) (Stm (Rep ADM))
 -> ADM (Stm (Rep ADM)))
-> BuilderT SOACS (ReaderT REnv (State RState)) (Stm (Rep ADM))
-> ADM (Stm (Rep ADM))
forall a b. (a -> b) -> a -> b
$ [VName]
-> Exp (Rep (BuilderT SOACS (ReaderT REnv (State RState))))
-> BuilderT
     SOACS
     (ReaderT REnv (State RState))
     (Stm (Rep (BuilderT SOACS (ReaderT REnv (State RState)))))
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m (Stm (Rep m))
mkLetNamesM [VName]
pat Exp (Rep (BuilderT SOACS (ReaderT REnv (State RState))))
Exp (Rep ADM)
e

  addStms :: Stms (Rep ADM) -> ADM ()
addStms = BuilderT SOACS (ReaderT REnv (State RState)) () -> ADM ()
forall a. BuilderT SOACS (ReaderT REnv (State RState)) a -> ADM a
ADM (BuilderT SOACS (ReaderT REnv (State RState)) () -> ADM ())
-> (Stms SOACS -> BuilderT SOACS (ReaderT REnv (State RState)) ())
-> Stms SOACS
-> ADM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms (Rep (BuilderT SOACS (ReaderT REnv (State RState))))
-> BuilderT SOACS (ReaderT REnv (State RState)) ()
Stms SOACS -> BuilderT SOACS (ReaderT REnv (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 REnv (State RState)) a
m) = BuilderT SOACS (ReaderT REnv (State RState)) (a, Stms (Rep ADM))
-> ADM (a, Stms (Rep ADM))
forall a. BuilderT SOACS (ReaderT REnv (State RState)) a -> ADM a
ADM (BuilderT SOACS (ReaderT REnv (State RState)) (a, Stms (Rep ADM))
 -> ADM (a, Stms (Rep ADM)))
-> BuilderT SOACS (ReaderT REnv (State RState)) (a, Stms (Rep ADM))
-> ADM (a, Stms (Rep ADM))
forall a b. (a -> b) -> a -> b
$ BuilderT SOACS (ReaderT REnv (State RState)) a
-> BuilderT
     SOACS
     (ReaderT REnv (State RState))
     (a, Stms (Rep (BuilderT SOACS (ReaderT REnv (State RState)))))
forall a.
BuilderT SOACS (ReaderT REnv (State RState)) a
-> BuilderT
     SOACS
     (ReaderT REnv (State RState))
     (a, Stms (Rep (BuilderT SOACS (ReaderT REnv (State RState)))))
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms BuilderT SOACS (ReaderT REnv (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 REnv (State RState)) Shape -> ADM Shape
forall a. BuilderT SOACS (ReaderT REnv (State RState)) a -> ADM a
ADM (BuilderT SOACS (ReaderT REnv (State RState)) Shape -> ADM Shape)
-> BuilderT SOACS (ReaderT REnv (State RState)) Shape -> ADM Shape
forall a b. (a -> b) -> a -> b
$ ReaderT REnv (State RState) Shape
-> BuilderT SOACS (ReaderT REnv (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 REnv (State RState) Shape
 -> BuilderT SOACS (ReaderT REnv (State RState)) Shape)
-> ReaderT REnv (State RState) Shape
-> BuilderT SOACS (ReaderT REnv (State RState)) Shape
forall a b. (a -> b) -> a -> b
$ (REnv -> Shape) -> ReaderT REnv (State RState) Shape
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks REnv -> Shape
envAdjShape

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 REnv (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 REnv (State RState) a -> REnv -> 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 REnv (State RState) (a, Stms SOACS)
-> ReaderT REnv (State RState) a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BuilderT SOACS (ReaderT REnv (State RState)) a
-> Scope SOACS -> ReaderT REnv (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 REnv (State RState)) a
m Scope SOACS
forall a. Monoid a => a
mempty) (REnv -> State RState a) -> REnv -> State RState a
forall a b. (a -> b) -> a -> b
$
            Shape -> Attrs -> Names -> REnv
REnv Shape
shape Attrs
attrs Names
forall a. Monoid a => a
mempty
        )
        (Map VName Adj
-> Substitutions -> Substitutions -> VNameSource -> RState
RState Map VName Adj
forall a. Monoid a => a
mempty Substitutions
forall a. Monoid a => a
mempty Substitutions
forall a. Monoid a => a
mempty VNameSource
vn)

adjVal :: Adj -> ADM VName
adjVal :: Adj -> ADM VName
adjVal (AdjVal SubExp
se) = Name -> Exp (Rep ADM) -> ADM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"const_val_adj" (Exp (Rep ADM) -> ADM VName) -> Exp (Rep ADM) -> ADM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep ADM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep ADM)) -> BasicOp -> Exp (Rep ADM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
adjVal (AdjSparse Sparse
sparse) = Sparse -> ADM VName
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ SOACS) =>
Sparse -> m VName
sparseArray Sparse
sparse
adjVal (AdjZero Shape
shape PrimType
t) = Shape -> TypeBase Shape NoMode -> ADM VName
forall (m :: * -> *).
MonadBuilder m =>
Shape -> TypeBase Shape NoMode -> m VName
zeroArray Shape
shape (TypeBase Shape NoMode -> ADM VName)
-> TypeBase Shape NoMode -> ADM VName
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
t

-- | Set a specific adjoint.
setAdj :: VName -> Adj -> ADM ()
setAdj :: VName -> Adj -> ADM ()
setAdj VName
v Adj
v_adj = (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 {stateAdjs = M.insert v v_adj $ stateAdjs env}

-- | Set an 'AdjVal' adjoint.  Simple wrapper around 'setAdj'.
insAdj :: VName -> VName -> ADM ()
insAdj :: VName -> VName -> ADM ()
insAdj VName
v = VName -> Adj -> ADM ()
setAdj VName
v (Adj -> ADM ()) -> (VName -> Adj) -> VName -> ADM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> Adj
AdjVal (SubExp -> Adj) -> (VName -> SubExp) -> VName -> Adj
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var

adjVName :: VName -> ADM VName
adjVName :: VName -> ADM VName
adjVName 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
"_adj")

-- | Create copies of all arrays consumed in the given statement, and
-- return statements which include copies of the consumed arrays.
--
-- See Note [Consumption].
copyConsumedArrsInStm :: Stm SOACS -> ADM (Substitutions, Stms SOACS)
copyConsumedArrsInStm :: Stm SOACS -> ADM (Substitutions, Stms SOACS)
copyConsumedArrsInStm Stm SOACS
s = Stm SOACS
-> ADM (Substitutions, Stms (Rep ADM))
-> ADM (Substitutions, Stms (Rep ADM))
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Stm SOACS
s (ADM (Substitutions, Stms (Rep ADM))
 -> ADM (Substitutions, Stms (Rep ADM)))
-> ADM (Substitutions, Stms (Rep ADM))
-> ADM (Substitutions, Stms (Rep ADM))
forall a b. (a -> b) -> a -> b
$ ADM Substitutions -> ADM (Substitutions, 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 Substitutions -> ADM (Substitutions, Stms (Rep ADM)))
-> ADM Substitutions -> ADM (Substitutions, Stms (Rep ADM))
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> ADM Substitutions
copyConsumedArrsInStm' Stm SOACS
s
  where
    copyConsumedArrsInStm' :: Stm SOACS -> ADM Substitutions
copyConsumedArrsInStm' Stm SOACS
stm =
      let onConsumed :: VName -> ADM [(VName, VName)]
onConsumed VName
v = Stm SOACS -> ADM [(VName, VName)] -> ADM [(VName, VName)]
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Stm SOACS
s (ADM [(VName, VName)] -> ADM [(VName, VName)])
-> ADM [(VName, VName)] -> ADM [(VName, VName)]
forall a b. (a -> b) -> a -> b
$ do
            v_t <- VName -> ADM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
v
            case v_t of
              Array {} -> do
                v' <-
                  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
"_ad_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)
                addSubstitution v' v
                pure [(v, v')]
              TypeBase Shape NoMode
_ -> [(VName, VName)] -> ADM [(VName, VName)]
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [(VName, VName)]
forall a. Monoid a => a
mempty

          consumed :: [VName]
consumed =
            Names -> [VName]
namesToList (Names -> [VName])
-> (Stms (Aliases SOACS) -> Names)
-> Stms (Aliases SOACS)
-> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms (Aliases SOACS) -> Names
forall rep. Aliased rep => Stms rep -> Names
consumedInStms (Stms (Aliases SOACS) -> [VName])
-> Stms (Aliases SOACS) -> [VName]
forall a b. (a -> b) -> a -> b
$
              (Stms (Aliases SOACS), AliasesAndConsumed) -> Stms (Aliases SOACS)
forall a b. (a, b) -> a
fst (AliasTable
-> Stms SOACS -> (Stms (Aliases SOACS), AliasesAndConsumed)
forall rep.
AliasableRep rep =>
AliasTable -> Stms rep -> (Stms (Aliases rep), AliasesAndConsumed)
Alias.analyseStms AliasTable
forall a. Monoid a => a
mempty (Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
stm))
       in [(VName, VName)] -> Substitutions
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, VName)] -> Substitutions)
-> ([[(VName, VName)]] -> [(VName, VName)])
-> [[(VName, VName)]]
-> Substitutions
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[(VName, VName)]] -> [(VName, VName)]
forall a. Monoid a => [a] -> a
mconcat ([[(VName, VName)]] -> Substitutions)
-> ADM [[(VName, VName)]] -> ADM Substitutions
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> ADM [(VName, VName)])
-> [VName] -> ADM [[(VName, 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, VName)]
onConsumed [VName]
consumed

copyConsumedArrsInBody :: [VName] -> Body SOACS -> ADM Substitutions
copyConsumedArrsInBody :: [VName] -> Body SOACS -> ADM Substitutions
copyConsumedArrsInBody [VName]
dontCopy Body SOACS
b =
  [Substitutions] -> Substitutions
forall a. Monoid a => [a] -> a
mconcat ([Substitutions] -> Substitutions)
-> ADM [Substitutions] -> ADM Substitutions
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> ADM Substitutions) -> [VName] -> ADM [Substitutions]
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 Substitutions
forall {m :: * -> *}. MonadBuilder m => VName -> m Substitutions
onConsumed ((VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [VName]
dontCopy) ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ GBody (Aliases SOACS) SubExpRes -> Names
forall res. GBody (Aliases SOACS) res -> Names
forall rep res. Aliased rep => GBody rep res -> Names
consumedInBody (AliasTable -> Body SOACS -> GBody (Aliases SOACS) SubExpRes
forall rep res.
(AliasableRep rep, IsResult res) =>
AliasTable -> GBody rep res -> GBody (Aliases rep) res
Alias.analyseBody AliasTable
forall a. Monoid a => a
mempty Body SOACS
b))
  where
    onConsumed :: VName -> m Substitutions
onConsumed VName
v = do
      v_t <- VName -> m (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
v
      case v_t of
        Acc {} -> [Char] -> m Substitutions
forall a. HasCallStack => [Char] -> a
error ([Char] -> m Substitutions) -> [Char] -> m Substitutions
forall a b. (a -> b) -> a -> b
$ [Char]
"copyConsumedArrsInBody: Acc " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString VName
v
        Array {} ->
          VName -> VName -> Substitutions
forall k a. k -> a -> Map k a
M.singleton VName
v
            (VName -> Substitutions) -> m VName -> m Substitutions
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Exp (Rep m) -> m 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
"_ad_copy")
              (BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ Shape -> SubExp -> BasicOp
Replicate Shape
forall a. Monoid a => a
mempty (VName -> SubExp
Var VName
v))
        TypeBase Shape NoMode
_ -> Substitutions -> m Substitutions
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Substitutions
forall a. Monoid a => a
mempty

returnSweepCode :: ADM a -> ADM a
returnSweepCode :: forall a. ADM a -> ADM a
returnSweepCode ADM a
m = do
  (a, stms) <- ADM a -> ADM (a, 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 a
m
  substs <- gets stateSubsts
  addStms $ substituteNames substs stms
  pure a

addSubstitution :: VName -> VName -> ADM ()
addSubstitution :: VName -> VName -> ADM ()
addSubstitution 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 {stateSubsts = M.insert v v' $ stateSubsts env}

-- While evaluating this action, pretend these variables have no
-- adjoints.  Restore current adjoints afterwards.  This is used for
-- handling certain nested operations. XXX: feels like this should
-- really be part of subAD, somehow.  Main challenge is that we don't
-- want to blank out Accumulator adjoints.  Also, might be inefficient
-- to blank out array adjoints.
noAdjsFor :: Names -> ADM a -> ADM a
noAdjsFor :: forall a. Names -> ADM a -> ADM a
noAdjsFor Names
names ADM a
m = do
  old <- (RState -> [Adj]) -> ADM [Adj]
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((RState -> [Adj]) -> ADM [Adj]) -> (RState -> [Adj]) -> ADM [Adj]
forall a b. (a -> b) -> a -> b
$ \RState
env -> (VName -> Maybe Adj) -> [VName] -> [Adj]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (VName -> Map VName Adj -> Maybe Adj
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` RState -> Map VName Adj
stateAdjs RState
env) [VName]
names'
  modify $ \RState
env -> RState
env {stateAdjs = foldl' (flip M.delete) (stateAdjs env) names'}
  x <- m
  modify $ \RState
env -> RState
env {stateAdjs = M.fromList (zip names' old) <> stateAdjs env}
  pure x
  where
    names' :: [VName]
names' = Names -> [VName]
namesToList Names
names

tabNest ::
  (MonadBuilder m, Rep m ~ SOACS) =>
  Int ->
  [VName] ->
  ([VName] -> [VName] -> m [VName]) ->
  m [VName]
tabNest :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ SOACS) =>
Int -> [VName] -> ([VName] -> [VName] -> m [VName]) -> m [VName]
tabNest = [VName]
-> Int -> [VName] -> ([VName] -> [VName] -> m [VName]) -> m [VName]
forall {m :: * -> *} {t}.
(BodyDec (Rep m) ~ (), OpC (Rep m) ~ SOAC, Eq t, Num t,
 MonadBuilder m, Buildable (Rep m)) =>
[VName]
-> t -> [VName] -> ([VName] -> [VName] -> m [VName]) -> m [VName]
tabNest' []
  where
    tabNest' :: [VName]
-> t -> [VName] -> ([VName] -> [VName] -> m [VName]) -> m [VName]
tabNest' [VName]
is t
0 [VName]
vs [VName] -> [VName] -> m [VName]
f = [VName] -> [VName] -> m [VName]
f ([VName] -> [VName]
forall a. [a] -> [a]
reverse [VName]
is) [VName]
vs
    tabNest' [VName]
is t
n [VName]
vs [VName] -> [VName] -> m [VName]
f = do
      vs_ts <- (VName -> m (TypeBase Shape NoMode))
-> [VName] -> m [TypeBase Shape NoMode]
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 -> m (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType [VName]
vs
      let w = Int -> [TypeBase Shape NoMode] -> SubExp
forall o. Int -> [TypeBase Shape o] -> SubExp
arraysSize Int
0 [TypeBase Shape NoMode]
vs_ts
      iota <-
        letExp "tab_iota" . BasicOp $
          Iota w (intConst Int64 0) (intConst Int64 1) Int64
      iparam <- newParam "i" $ Prim int64
      params <- forM vs $ \VName
v ->
        Name -> TypeBase Shape NoMode -> m (Param (TypeBase Shape NoMode))
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam (VName -> Name
baseName VName
v Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_p") (TypeBase Shape NoMode -> m (Param (TypeBase Shape NoMode)))
-> (TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> TypeBase Shape NoMode
-> m (Param (TypeBase Shape NoMode))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Shape NoMode -> TypeBase Shape NoMode
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType (TypeBase Shape NoMode -> m (Param (TypeBase Shape NoMode)))
-> m (TypeBase Shape NoMode) -> m (Param (TypeBase Shape NoMode))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> m (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
v
      ((ret, res), stms) <- collectStms . localScope (scopeOfLParams (iparam : params)) $ do
        res <- tabNest' (paramName iparam : is) (n - 1) (map paramName params) f
        ret <- mapM lookupType res
        pure (ret, varsRes res)
      let lam = [LParam (Rep m)]
-> [TypeBase Shape NoMode]
-> Body (Rep m)
-> GLambda (Rep m) (TypeBase Shape NoMode)
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda (Param (TypeBase Shape NoMode)
iparam Param (TypeBase Shape NoMode)
-> [Param (TypeBase Shape NoMode)]
-> [Param (TypeBase Shape NoMode)]
forall a. a -> [a] -> [a]
: [Param (TypeBase Shape NoMode)]
params) [TypeBase Shape NoMode]
ret (BodyDec (Rep m) -> Stms (Rep m) -> [SubExpRes] -> Body (Rep m)
forall rep res. BodyDec rep -> Stms rep -> [res] -> GBody rep res
Body () Stms (Rep m)
stms [SubExpRes]
res)
      letTupExp "tab" . Op . Screma w (iota : vs) =<< mapSOAC lam

-- | Construct a lambda for binop'ing two values of the given type,
-- which may be arrays.
vecOpLambda :: (PrimType -> BinOp) -> Type -> ADM (Lambda SOACS)
vecOpLambda :: (PrimType -> BinOp) -> TypeBase Shape NoMode -> ADM (Lambda SOACS)
vecOpLambda PrimType -> BinOp
bop (Prim PrimType
pt) = BinOp -> PrimType -> ADM (Lambda (Rep ADM))
forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m)) =>
BinOp -> PrimType -> m (Lambda (Rep m))
binOpLambda (PrimType -> BinOp
bop PrimType
pt) PrimType
pt
vecOpLambda PrimType -> BinOp
bop t :: TypeBase Shape NoMode
t@Array {} = do
  xs_p <- Name
-> TypeBase Shape NoMode -> ADM (Param (TypeBase Shape NoMode))
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"xs" TypeBase Shape NoMode
t
  ys_p <- newParam "ys" t
  lam <- vecOpLambda bop $ rowType t
  body <- insertStmsM $ do
    res <-
      letSubExp "lam_map"
        . Op
        . Screma (arraySize 0 t) [paramName xs_p, paramName ys_p]
        =<< mapSOAC lam
    pure $ resultBody [res]
  pure
    Lambda
      { lambdaParams = [xs_p, ys_p],
        lambdaReturnType = [t],
        lambdaBody = body
      }
vecOpLambda PrimType -> BinOp
_ TypeBase Shape NoMode
t =
  [Char] -> ADM (Lambda SOACS)
forall a. HasCallStack => [Char] -> a
error ([Char] -> ADM (Lambda SOACS)) -> [Char] -> ADM (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ [Char]
"vecOpLambda: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ TypeBase Shape NoMode -> [Char]
forall a. Show a => a -> [Char]
show TypeBase Shape NoMode
t

-- | Construct an expression for adding the two variables.
addExp :: VName -> VName -> ADM (Exp SOACS)
addExp :: VName -> VName -> ADM (Exp SOACS)
addExp VName
x VName
y = do
  x_t <- VName -> ADM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
x
  case x_t of
    Prim PrimType
pt ->
      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
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp (PrimType -> BinOp
addBinOp PrimType
pt) (VName -> SubExp
Var VName
x) (VName -> SubExp
Var VName
y)
    Array {} -> do
      lam <- TypeBase Shape NoMode -> ADM (Lambda (Rep ADM))
forall (m :: * -> *).
(OpC (Rep m) ~ SOAC, MonadBuilder m, Buildable (Rep m)) =>
TypeBase Shape NoMode -> m (Lambda (Rep m))
addLambda (TypeBase Shape NoMode -> ADM (Lambda (Rep ADM)))
-> TypeBase Shape NoMode -> ADM (Lambda (Rep ADM))
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoMode -> TypeBase Shape NoMode
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType TypeBase Shape NoMode
x_t
      Op . Screma (arraySize 0 x_t) [x, y] <$> mapSOAC lam
    TypeBase Shape NoMode
_ ->
      [Char] -> ADM (Exp SOACS)
forall a. HasCallStack => [Char] -> a
error ([Char] -> ADM (Exp SOACS)) -> [Char] -> ADM (Exp SOACS)
forall a b. (a -> b) -> a -> b
$ [Char]
"addExp: unexpected type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ TypeBase Shape NoMode -> [Char]
forall a. Pretty a => a -> [Char]
prettyString TypeBase Shape NoMode
x_t

-- | Construct an expression for performing this binary operation on two variables.
vecOpExp :: (PrimType -> BinOp) -> VName -> VName -> ADM (Exp SOACS)
vecOpExp :: (PrimType -> BinOp) -> VName -> VName -> ADM (Exp SOACS)
vecOpExp PrimType -> BinOp
bop VName
x VName
y = do
  x_t <- VName -> ADM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
x
  case x_t of
    Prim PrimType
pt ->
      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
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp (PrimType -> BinOp
bop PrimType
pt) (VName -> SubExp
Var VName
x) (VName -> SubExp
Var VName
y)
    Array {} -> do
      lam <- (PrimType -> BinOp) -> TypeBase Shape NoMode -> ADM (Lambda SOACS)
vecOpLambda PrimType -> BinOp
bop (TypeBase Shape NoMode -> ADM (Lambda SOACS))
-> TypeBase Shape NoMode -> ADM (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoMode -> TypeBase Shape NoMode
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType TypeBase Shape NoMode
x_t
      Op . Screma (arraySize 0 x_t) [x, y] <$> mapSOAC lam
    TypeBase Shape NoMode
_ ->
      [Char] -> ADM (Exp SOACS)
forall a. HasCallStack => [Char] -> a
error ([Char] -> ADM (Exp SOACS)) -> [Char] -> ADM (Exp SOACS)
forall a b. (a -> b) -> a -> b
$ [Char]
"vecOpExp: unexpected type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ TypeBase Shape NoMode -> [Char]
forall a. Pretty a => a -> [Char]
prettyString TypeBase Shape NoMode
x_t

lookupAdj :: VName -> ADM Adj
lookupAdj :: VName -> ADM Adj
lookupAdj VName
v = do
  maybeAdj <- (RState -> Maybe Adj) -> ADM (Maybe Adj)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((RState -> Maybe Adj) -> ADM (Maybe Adj))
-> (RState -> Maybe Adj) -> ADM (Maybe Adj)
forall a b. (a -> b) -> a -> b
$ VName -> Map VName Adj -> Maybe Adj
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (Map VName Adj -> Maybe Adj)
-> (RState -> Map VName Adj) -> RState -> Maybe Adj
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RState -> Map VName Adj
stateAdjs
  adj_shape <- askShape
  case maybeAdj of
    Maybe Adj
Nothing -> do
      v_t <- VName -> ADM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
v
      case v_t of
        Acc VName
_ Shape
shape [Prim PrimType
t] -> Adj -> ADM Adj
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Adj -> ADM Adj) -> Adj -> ADM Adj
forall a b. (a -> b) -> a -> b
$ Shape -> PrimType -> Adj
AdjZero Shape
shape PrimType
t
        Acc VName
_ Shape
shape [TypeBase Shape NoMode
t] -> Adj -> ADM Adj
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Adj -> ADM Adj) -> Adj -> ADM Adj
forall a b. (a -> b) -> a -> b
$ Shape -> PrimType -> Adj
AdjZero (Shape
shape Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> TypeBase Shape NoMode -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape TypeBase Shape NoMode
t) (TypeBase Shape NoMode -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType TypeBase Shape NoMode
t)
        Acc {} -> [Char] -> ADM Adj
forall a. HasCallStack => [Char] -> a
error ([Char] -> ADM Adj) -> [Char] -> ADM Adj
forall a b. (a -> b) -> a -> b
$ [Char]
"lookupAdj: Non-singleton accumulator adjoint: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> TypeBase Shape NoMode -> [Char]
forall a. Pretty a => a -> [Char]
prettyString TypeBase Shape NoMode
v_t
        TypeBase Shape NoMode
_ -> Adj -> ADM Adj
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Adj -> ADM Adj) -> Adj -> ADM Adj
forall a b. (a -> b) -> a -> b
$ Shape -> PrimType -> Adj
AdjZero (Shape
adj_shape Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> TypeBase Shape NoMode -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape TypeBase Shape NoMode
v_t) (TypeBase Shape NoMode -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType TypeBase Shape NoMode
v_t)
    Just Adj
v_adj -> Adj -> ADM Adj
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Adj
v_adj

lookupAdjVal :: VName -> ADM VName
lookupAdjVal :: VName -> ADM VName
lookupAdjVal VName
v = Adj -> ADM VName
adjVal (Adj -> ADM VName) -> ADM Adj -> ADM VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> ADM Adj
lookupAdj VName
v

updateAdjIndex :: VName -> (InBounds, SubExp) -> SubExp -> ADM ()
updateAdjIndex :: VName -> (InBounds, SubExp) -> SubExp -> ADM ()
updateAdjIndex VName
v (InBounds
check, SubExp
i) SubExp
se = do
  maybeAdj <- (RState -> Maybe Adj) -> ADM (Maybe Adj)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((RState -> Maybe Adj) -> ADM (Maybe Adj))
-> (RState -> Maybe Adj) -> ADM (Maybe Adj)
forall a b. (a -> b) -> a -> b
$ VName -> Map VName Adj -> Maybe Adj
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (Map VName Adj -> Maybe Adj)
-> (RState -> Map VName Adj) -> RState -> Maybe Adj
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RState -> Map VName Adj
stateAdjs
  t <- lookupType v
  adj_shape <- askShape
  let iv = (InBounds
check, SubExp
i, SubExp
se)
      vec_dims = Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
adj_shape
      full_shape = Shape
adj_shape Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> TypeBase Shape NoMode -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape TypeBase Shape NoMode
t
  case maybeAdj of
    Maybe Adj
Nothing ->
      VName -> Adj -> ADM ()
setAdj VName
v (Adj -> ADM ()) -> Adj -> ADM ()
forall a b. (a -> b) -> a -> b
$ Sparse -> Adj
AdjSparse (Sparse -> Adj) -> Sparse -> Adj
forall a b. (a -> b) -> a -> b
$ Shape -> PrimType -> Int -> [(InBounds, SubExp, SubExp)] -> Sparse
Sparse Shape
full_shape (TypeBase Shape NoMode -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType TypeBase Shape NoMode
t) Int
vec_dims [(InBounds, SubExp, SubExp)
iv]
    Just (AdjZero {}) ->
      VName -> Adj -> ADM ()
setAdj VName
v (Adj -> ADM ()) -> Adj -> ADM ()
forall a b. (a -> b) -> a -> b
$ Sparse -> Adj
AdjSparse (Sparse -> Adj) -> Sparse -> Adj
forall a b. (a -> b) -> a -> b
$ Shape -> PrimType -> Int -> [(InBounds, SubExp, SubExp)] -> Sparse
Sparse Shape
full_shape (TypeBase Shape NoMode -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType TypeBase Shape NoMode
t) Int
vec_dims [(InBounds, SubExp, SubExp)
iv]
    Just (AdjSparse (Sparse Shape
shape PrimType
pt Int
vd [(InBounds, SubExp, SubExp)]
ivs)) ->
      VName -> Adj -> ADM ()
setAdj VName
v (Adj -> ADM ()) -> Adj -> ADM ()
forall a b. (a -> b) -> a -> b
$ Sparse -> Adj
AdjSparse (Sparse -> Adj) -> Sparse -> Adj
forall a b. (a -> b) -> a -> b
$ Shape -> PrimType -> Int -> [(InBounds, SubExp, SubExp)] -> Sparse
Sparse Shape
shape PrimType
pt Int
vd ([(InBounds, SubExp, SubExp)] -> Sparse)
-> [(InBounds, SubExp, SubExp)] -> Sparse
forall a b. (a -> b) -> a -> b
$ (InBounds, SubExp, SubExp)
iv (InBounds, SubExp, SubExp)
-> [(InBounds, SubExp, SubExp)] -> [(InBounds, SubExp, SubExp)]
forall a. a -> [a] -> [a]
: [(InBounds, SubExp, SubExp)]
ivs
    Just adj :: Adj
adj@AdjVal {} -> do
      v_adj <- Adj -> ADM VName
adjVal Adj
adj
      v_adj_t <- lookupType v_adj
      se_v <- letExp "se_v" $ BasicOp $ SubExp se
      vec_shape <- askShape
      insAdj v
        =<< case v_adj_t of
          Acc {} -> do
            let stms :: Safety -> ADM VName
stms Safety
s = do
                  attrs <- (REnv -> Attrs) -> ADM Attrs
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks REnv -> Attrs
envAttrs
                  dims <- arrayDims <$> lookupType se_v
                  ~[v_adj'] <-
                    attributing attrs $
                      tabNest (length dims) [se_v, v_adj] $ \[VName]
is [VName
se_v', VName
v_adj'] -> do
                        let ([SubExp]
vec_is, [SubExp]
val_is) = Int -> [SubExp] -> ([SubExp], [SubExp])
forall a. Int -> [a] -> ([a], [a])
splitAt (Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
vec_shape) ([SubExp] -> ([SubExp], [SubExp]))
-> [SubExp] -> ([SubExp], [SubExp])
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
is
                        Name -> Exp (Rep ADM) -> ADM [VName]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [VName]
letTupExp Name
"acc" (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
$
                          Safety -> VName -> [SubExp] -> [SubExp] -> BasicOp
UpdateAcc Safety
s VName
v_adj' ([SubExp]
vec_is [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ SubExp
i SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: [SubExp]
val_is) [VName -> SubExp
Var VName
se_v']
                  pure v_adj'
            case InBounds
check of
              CheckBounds Maybe SubExp
_ -> Safety -> ADM VName
stms Safety
Safe
              InBounds
AssumeBounds -> Safety -> ADM VName
stms Safety
Unsafe
              InBounds
OutOfBounds -> VName -> ADM VName
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v_adj
          TypeBase Shape NoMode
_ -> do
            let stms :: Safety -> ADM VName
stms Safety
s = do
                  let slice :: Slice SubExp
slice =
                        TypeBase Shape NoMode -> [DimIndex SubExp] -> Slice SubExp
fullSlice TypeBase Shape NoMode
v_adj_t ([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
vec_shape) [DimIndex SubExp] -> [DimIndex SubExp] -> [DimIndex SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix SubExp
i]
                  v_adj_i <-
                    Name -> Exp (Rep ADM) -> ADM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp (VName -> Name
baseName VName
v_adj Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_i") (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 -> Slice SubExp -> BasicOp
Index VName
v_adj Slice SubExp
slice
                  se_update <- letSubExp "updated_adj_i" =<< addExp se_v v_adj_i
                  letExp (baseName v_adj) . BasicOp $
                    Update s v_adj slice se_update
            case InBounds
check of
              CheckBounds Maybe SubExp
_ -> Safety -> ADM VName
stms Safety
Safe
              InBounds
AssumeBounds -> Safety -> ADM VName
stms Safety
Unsafe
              InBounds
OutOfBounds -> VName -> ADM VName
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v_adj

updateAdjWithSafety :: VName -> VName -> Safety -> ADM ()
updateAdjWithSafety :: VName -> VName -> Safety -> ADM ()
updateAdjWithSafety VName
v VName
d Safety
safety = do
  maybeAdj <- (RState -> Maybe Adj) -> ADM (Maybe Adj)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((RState -> Maybe Adj) -> ADM (Maybe Adj))
-> (RState -> Maybe Adj) -> ADM (Maybe Adj)
forall a b. (a -> b) -> a -> b
$ VName -> Map VName Adj -> Maybe Adj
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (Map VName Adj -> Maybe Adj)
-> (RState -> Map VName Adj) -> RState -> Maybe Adj
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RState -> Map VName Adj
stateAdjs
  case maybeAdj of
    Maybe Adj
Nothing ->
      VName -> VName -> ADM ()
insAdj VName
v VName
d
    Just Adj
adj -> do
      v_adj <- Adj -> ADM VName
adjVal Adj
adj
      v_adj_t <- lookupType v_adj
      case v_adj_t of
        Acc {} -> do
          dims <- TypeBase Shape NoMode -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims (TypeBase Shape NoMode -> [SubExp])
-> ADM (TypeBase Shape NoMode) -> ADM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> ADM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
d
          ~[v_adj'] <-
            tabNest (length dims) [d, v_adj] $ \[VName]
is [VName
d', VName
v_adj'] ->
              Name -> Exp (Rep ADM) -> ADM [VName]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [VName]
letTupExp Name
"acc" (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
$
                Safety -> VName -> [SubExp] -> [SubExp] -> BasicOp
UpdateAcc Safety
safety VName
v_adj' ((VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
is) [VName -> SubExp
Var VName
d']
          insAdj v v_adj'
        TypeBase Shape NoMode
_ -> do
          v_adj' <- 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
"_adj") (Exp SOACS -> ADM VName) -> ADM (Exp SOACS) -> ADM VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> VName -> ADM (Exp SOACS)
addExp VName
v_adj VName
d
          insAdj v v_adj'

updateAdjSliceWithSafety :: Slice SubExp -> VName -> VName -> Safety -> ADM ()
updateAdjSliceWithSafety :: Slice SubExp -> VName -> VName -> Safety -> ADM ()
updateAdjSliceWithSafety (Slice [DimFix SubExp
i]) VName
v VName
d Safety
safety =
  VName -> (InBounds, SubExp) -> SubExp -> ADM ()
updateAdjIndex VName
v (InBounds
bounds, SubExp
i) (VName -> SubExp
Var VName
d)
  where
    bounds :: InBounds
bounds = case Safety
safety of
      Safety
Safe -> Maybe SubExp -> InBounds
CheckBounds Maybe SubExp
forall a. Maybe a
Nothing
      Safety
Unsafe -> InBounds
AssumeBounds
updateAdjSliceWithSafety Slice SubExp
slice VName
v VName
d Safety
safety = do
  t <- VName -> ADM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
v
  v_adj <- lookupAdjVal v
  v_adj_t <- lookupType v_adj
  v_adj' <- case v_adj_t of
    Acc {} -> do
      let dims :: [SubExp]
dims = Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice
      ~[v_adj'] <-
        Int
-> [VName] -> ([VName] -> [VName] -> ADM [VName]) -> ADM [VName]
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ SOACS) =>
Int -> [VName] -> ([VName] -> [VName] -> m [VName]) -> m [VName]
tabNest ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
dims) [VName
d, VName
v_adj] (([VName] -> [VName] -> ADM [VName]) -> ADM [VName])
-> ([VName] -> [VName] -> ADM [VName]) -> ADM [VName]
forall a b. (a -> b) -> a -> b
$ \[VName]
is [VName
d', VName
v_adj'] -> do
          slice' <-
            (TPrimExp Int64 VName -> ADM SubExp)
-> [TPrimExp Int64 VName] -> ADM [SubExp]
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) -> [a] -> f [b]
traverse (Name -> TPrimExp Int64 VName -> ADM SubExp
forall (m :: * -> *) a.
(MonadBuilder m, ToExp a) =>
Name -> a -> m SubExp
toSubExp Name
"index") ([TPrimExp Int64 VName] -> ADM [SubExp])
-> [TPrimExp Int64 VName] -> ADM [SubExp]
forall a b. (a -> b) -> a -> b
$
              Slice (TPrimExp Int64 VName)
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall d. Num d => Slice d -> [d] -> [d]
fixSlice ((SubExp -> TPrimExp Int64 VName)
-> Slice SubExp -> Slice (TPrimExp Int64 VName)
forall a b. (a -> b) -> Slice a -> Slice b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 Slice SubExp
slice) ([TPrimExp Int64 VName] -> [TPrimExp Int64 VName])
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> a -> b
$
                (VName -> TPrimExp Int64 VName)
-> [VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 [VName]
is
          letTupExp (baseName v_adj') . BasicOp $
            UpdateAcc safety v_adj' slice' [Var d']
      pure v_adj'
    TypeBase Shape NoMode
_ -> do
      vec_slice <- TypeBase Shape NoMode -> [DimIndex SubExp] -> ADM (Slice SubExp)
vecSlice TypeBase Shape NoMode
v_adj_t ([DimIndex SubExp] -> ADM (Slice SubExp))
-> [DimIndex SubExp] -> ADM (Slice SubExp)
forall a b. (a -> b) -> a -> b
$ Slice SubExp -> [DimIndex SubExp]
forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slice
      v_adjslice <-
        if primType t
          then pure v_adj
          else letExp (baseName v <> "_slice") $ BasicOp $ Index v_adj vec_slice
      letInPlace "updated_adj" v_adj vec_slice =<< addExp v_adjslice d
  insAdj v v_adj'

updateAdj :: VName -> VName -> ADM ()
updateAdj :: VName -> VName -> ADM ()
updateAdj VName
v VName
d = VName -> VName -> Safety -> ADM ()
updateAdjWithSafety VName
v VName
d Safety
Unsafe

updateAdjSlice :: Slice SubExp -> VName -> VName -> ADM ()
updateAdjSlice :: Slice SubExp -> VName -> VName -> ADM ()
updateAdjSlice Slice SubExp
slice VName
v VName
d = Slice SubExp -> VName -> VName -> Safety -> ADM ()
updateAdjSliceWithSafety Slice SubExp
slice VName
v VName
d Safety
Unsafe

updateSubExpAdj :: SubExp -> VName -> ADM ()
updateSubExpAdj :: SubExp -> VName -> ADM ()
updateSubExpAdj Constant {} VName
_ = () -> ADM ()
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
updateSubExpAdj (Var VName
v) VName
d = 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
d

-- | Is this primal variable active in the AD sense?  FIXME: this is
-- (obviously) much too conservative.
isActive :: VName -> ADM Bool
isActive :: VName -> ADM Bool
isActive = (TypeBase Shape NoMode -> Bool)
-> ADM (TypeBase Shape NoMode) -> ADM Bool
forall a b. (a -> b) -> ADM a -> ADM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (TypeBase Shape NoMode -> TypeBase Shape NoMode -> Bool
forall a. Eq a => a -> a -> Bool
/= PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Unit) (ADM (TypeBase Shape NoMode) -> ADM Bool)
-> (VName -> ADM (TypeBase Shape NoMode)) -> VName -> ADM Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> ADM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType

-- | Ignore any changes to adjoints made while evaluating this action.
subAD :: ADM a -> ADM a
subAD :: forall a. ADM a -> ADM a
subAD ADM a
m = do
  old_state_adjs <- (RState -> Map VName Adj) -> ADM (Map VName Adj)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets RState -> Map VName Adj
stateAdjs
  x <- m
  modify $ \RState
s -> RState
s {stateAdjs = old_state_adjs}
  pure x

subSubsts :: ADM a -> ADM a
subSubsts :: forall a. ADM a -> ADM a
subSubsts ADM a
m = do
  old_state_substs <- (RState -> Substitutions) -> ADM Substitutions
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets RState -> Substitutions
stateSubsts
  x <- m
  modify $ \RState
s -> RState
s {stateSubsts = old_state_substs}
  pure x

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

-- | @setLoopTape v vs@ establishes @vs@ as the name of the array
-- where values of loop parameter @v@ from the forward pass are
-- stored.
setLoopTape :: VName -> VName -> ADM ()
setLoopTape :: VName -> VName -> ADM ()
setLoopTape VName
v VName
vs = (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 {stateLoopTape = M.insert v vs $ stateLoopTape env}

-- | Look-up the name of the array where @v@ is stored.
lookupLoopTape :: VName -> ADM (Maybe VName)
lookupLoopTape :: VName -> ADM (Maybe VName)
lookupLoopTape VName
v = (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 -> Substitutions -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (Substitutions -> Maybe VName)
-> (RState -> Substitutions) -> RState -> Maybe VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RState -> Substitutions
stateLoopTape

-- | @substLoopTape v v'@ substitutes the key @v@ for @v'@. That is,
-- if @v |-> vs@ then after the substitution @v' |-> vs@ (and @v@
-- points to nothing).
substLoopTape :: VName -> VName -> ADM ()
substLoopTape :: VName -> VName -> ADM ()
substLoopTape VName
v VName
v' = (VName -> ADM ()) -> Maybe VName -> ADM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (VName -> VName -> ADM ()
setLoopTape VName
v') (Maybe VName -> ADM ()) -> ADM (Maybe VName) -> ADM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> ADM (Maybe VName)
lookupLoopTape VName
v

-- | Renames the keys of the loop tape. Useful for fixing the
-- the names in the loop tape after a loop rename.
renameLoopTape :: Substitutions -> ADM ()
renameLoopTape :: Substitutions -> ADM ()
renameLoopTape = ((VName, VName) -> ADM ()) -> [(VName, VName)] -> ADM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((VName -> VName -> ADM ()) -> (VName, VName) -> ADM ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry VName -> VName -> ADM ()
substLoopTape) ([(VName, VName)] -> ADM ())
-> (Substitutions -> [(VName, VName)]) -> Substitutions -> ADM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Substitutions -> [(VName, VName)]
forall k a. Map k a -> [(k, a)]
M.toList

-- | Disable vector AD within the provided action. This results in a map that
-- computes each adjoint explicitly, then assembles the resulting adjoint
-- vectors. This is useful for constructs (such as scans) where vector AD is
-- impractical or inefficient.
locallyNonvector ::
  (FreeIn e) =>
  -- | Something that represents all the free variables used in the action.
  -- Usually just an expression or statement.
  e ->
  ADM () ->
  ADM ()
locallyNonvector :: forall e. FreeIn e => e -> ADM () -> ADM ()
locallyNonvector e
e ADM ()
m = do
  adj_shape <- ADM Shape
askShape
  if adj_shape == mempty
    then m
    else do
      -- We map over all adjoints of free variables in 'e'. To avoid clutter, we
      -- only consider those that actually have known nonzero adjoints.
      e_adjs <- filterM knownAdjoint e_free
      e_adjs_vals <- mapM lookupAdjVal e_adjs
      e_free_adjs <- mkMap "nonvec_adj" adj_shape e_adjs_vals $ \[VName]
e_adjs_vals' -> do
        (VName -> VName -> ADM ()) -> [VName] -> [VName] -> ADM ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ VName -> VName -> ADM ()
insAdj [VName]
e_adjs [VName]
e_adjs_vals'
        (REnv -> REnv) -> ADM () -> ADM ()
forall a. (REnv -> REnv) -> ADM a -> ADM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (\REnv
env -> REnv
env {envAdjShape = mempty}) ADM ()
m
        (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]
e_free
      zipWithM_ insAdj e_free e_free_adjs
  where
    e_free :: [VName]
e_free = Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ e -> Names
forall a. FreeIn a => a -> Names
freeIn e
e
    knownAdjoint :: VName -> ADM Bool
knownAdjoint VName
v = do
      v_adj <- VName -> ADM Adj
lookupAdj VName
v
      pure $ case v_adj of
        AdjZero {} -> Bool
False
        Adj
_ -> Bool
True

-- | Note that these accumulator certificates belong to accumulators that have
-- a combining operator, for the duration of the action.
withOperatorAccs :: [VName] -> ADM a -> ADM a
withOperatorAccs :: forall a. [VName] -> ADM a -> ADM a
withOperatorAccs [VName]
certs = (REnv -> REnv) -> ADM a -> ADM a
forall a. (REnv -> REnv) -> ADM a -> ADM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((REnv -> REnv) -> ADM a -> ADM a)
-> (REnv -> REnv) -> ADM a -> ADM a
forall a b. (a -> b) -> a -> b
$ \REnv
env ->
  REnv
env {envOperatorAccs = namesFromList certs <> envOperatorAccs env}

-- | Does this accumulator certificate belong to an accumulator with a
-- combining operator, rather than a scatter-like one that overwrites?  See Note
-- [Adjoints of accumulators].
isOperatorAcc :: VName -> ADM Bool
isOperatorAcc :: VName -> ADM Bool
isOperatorAcc VName
cert = BuilderT SOACS (ReaderT REnv (State RState)) Bool -> ADM Bool
forall a. BuilderT SOACS (ReaderT REnv (State RState)) a -> ADM a
ADM (BuilderT SOACS (ReaderT REnv (State RState)) Bool -> ADM Bool)
-> BuilderT SOACS (ReaderT REnv (State RState)) Bool -> ADM Bool
forall a b. (a -> b) -> a -> b
$ ReaderT REnv (State RState) Bool
-> BuilderT SOACS (ReaderT REnv (State RState)) Bool
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 REnv (State RState) Bool
 -> BuilderT SOACS (ReaderT REnv (State RState)) Bool)
-> ReaderT REnv (State RState) Bool
-> BuilderT SOACS (ReaderT REnv (State RState)) Bool
forall a b. (a -> b) -> a -> b
$ (REnv -> Bool) -> ReaderT REnv (State RState) Bool
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((REnv -> Bool) -> ReaderT REnv (State RState) Bool)
-> (REnv -> Bool) -> ReaderT REnv (State RState) Bool
forall a b. (a -> b) -> a -> b
$ (VName
cert VName -> Names -> Bool
`nameIn`) (Names -> Bool) -> (REnv -> Names) -> REnv -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. REnv -> Names
envOperatorAccs

-- | @vecSlice t is@ indexes an adjoint of type @t@ with @is@, which is written
-- against the type of the corresponding primal value.  As the vector
-- dimensions are outermost in the adjoint, they must be skipped past.
vecSlice :: Type -> [DimIndex SubExp] -> ADM (Slice SubExp)
vecSlice :: TypeBase Shape NoMode -> [DimIndex SubExp] -> ADM (Slice SubExp)
vecSlice TypeBase Shape NoMode
t [DimIndex SubExp]
is = do
  adj_shape <- ADM Shape
askShape
  pure $ sliceAt t (shapeRank adj_shape) is

-- | If we are doing vector AD, apply 'vecPerm' to the array.
vecToInner :: VName -> ADM VName
vecToInner :: VName -> ADM VName
vecToInner VName
v = do
  adj_shape <- ADM Shape
askShape
  if adj_shape == mempty
    then pure v
    else do
      v_t <- lookupType v
      letExp (baseName v <> "_tr") . BasicOp . Rearrange v $
        vecPerm adj_shape v_t

-- Note [Consumption]
--
-- Parts of this transformation depends on duplicating computation.
-- This is a problem when a primal expression consumes arrays (via
-- e.g. Update).  For example, consider how we handle this conditional:
--
--   if b then ys with [0] = 0 else ys
--
-- This consumes the array 'ys', which means that when we later
-- generate code for the return sweep, we can no longer use 'ys'.
-- This is a problem, because when we call 'diffBody' on the branch
-- bodies, we'll keep the primal code (maybe it'll be removed by
-- simplification later - we cannot know).  A similar issue occurs for
-- SOACs.  Our solution is to make copies of all consumes arrays:
--
--  let ys_copy = copy ys
--
-- Then we generate code for the return sweep as normal, but replace _every
-- instance_ of 'ys' in the generated code with 'ys_copy'. This works because
-- Futhark does not have *semantic* in-place updates - any consumption violation
-- can be replaced with copies (on arrays, anyway).
--
-- If we are lucky, the uses of 'ys_copy' will be removed by
-- simplification, and there will be no overhead.  But even if not,
-- this is still (asymptotically) efficient because the array that is
-- being consumed must in any case have been produced within the code
-- that we are differentiating, so a copy is at most a scalar
-- overhead.  This is _not_ the case when loops are involved.
--
-- Also, the above only works for arrays, not accumulator variables.
-- Those will need some other mechanism.