-- | Differentiation related to accumulators in the input program.
module Futhark.AD.Rev.Acc
  ( diffWithAcc,
    diffUpdateAcc,
  )
where

-- Note [Adjoints of accumulators]
--
-- The general case of taking adjoints of WithAcc is tricky.  We make
-- some assumptions and lay down a basic design.
--
-- First, we assume that any WithAccs that occur in the program come from one of
-- these sources:
--
-- - A previous instance of VJP, which means we can rely on the operator having
--   a constant adjoint (it's addition as appropriate to the type).
--
-- - A scatter, meaning there is no operator.
--
-- This means we are in fact ignoring one potential source:
--
-- - reduce_by_index_stream, where the operator is whatever the programmer
--   wrote.
--
-- This is because we assume that an operator (if it exists) is addition,
-- meaning it has a constant adjoint. This is acceptable because
-- reduce_by_index_stream is not a real part of the language, but is exposed
-- solely for testing the WithAcc machinery.
--
-- These are distinguished by the presence of an operator, which matters because
-- only the scatter overwrites: 'update_acc' on an accumulator that has an
-- operator combines with the value already in the cell, so that value keeps its
-- full sensitivity. See 'isOperatorAcc'.
--
-- Second, the adjoint of an accumulator is an array of the same type
-- as the underlying array.  For example, the adjoint type of the
-- primal type 'acc(c, [n], {f64})' is '[n]f64'.  In principle the
-- adjoint of 'acc(c, [n], {f64,f32})' should be two arrays of type
-- '[]f64', '[]f32'.  Our current design assumes that adjoints are
-- single variables.  This is fixable.
--
-- In the return sweep, when inserting the with_acc, we still compute the
-- "original" accumulator result, but modified such that its initial value is
-- the adjoint of the result of the accumulator. We also modify the update_accs
-- of these accumulators to be with zero values. For a scatter this means that
-- the array that is produced will be equal to the adjoint of the result, except
-- for those places that have been updated, where it will be zero. This is
-- intuitively sensible - values that have been overwritten (and so do not
-- contribute to the result) should obviously have zero sensitivity.
--
-- When the accumulator has an operator, nothing is overwritten, so the array
-- must instead come out as the adjoint of the result in its entirety. We get
-- that by giving the accumulator of the return sweep addition as its operator,
-- which makes those writes of zeroes leave it alone.
--
-- # Adjoint of UpdateAcc
--
-- Consider primal code
--
--     update_acc(acc, i, v)
--
-- Interpreted as an imperative statement, this means
--
--     acc[i] ⊕= v
--
-- for some '⊕'.  Normally all the compiler knows of '⊕' is that it
-- is associative and commutative, but because we assume that all
-- accumulators are the result of previous AD transformations, we
-- can assume that '⊕' actually behaves like addition - that is, has
-- unit partial derivatives.  So the return sweep is
--
--     v_adj += acc_adj[i]
--
-- and the adjoint of the accumulator going in is the adjoint of the one coming
-- out, except for a scatter, where the cell that is overwritten is zeroed.
--
-- Further, we modify the primal code so that it becomes
--
--     update_acc(acc, i, 0)
--
-- for some appropriate notion of zero.
--
-- # Adjoint of Map
--
-- Suppose we have primal code
--
--   let acc' =
--     map (...) acc
--
-- where "acc : acc(c, [n], {f64})" and the width of the Map is "w".
-- Our normal transformation for Map input arrays is to similarly map
-- their adjoint, but clearly this doesn't work here because the
-- semantics of mapping an adjoint is an "implicit replicate".  So
-- when generating the return sweep we actually perform that
-- replication:
--
--   map (...) (replicate w acc_adj)
--
-- But what about the contributions to "acc'"?  Those we also have to
-- take special care of.  The result of the map itself is actually a
-- multidimensional array:
--
--   let acc_contribs =
--     map (...) (replicate w acc'_adj)
--
-- which we must then sum to add to the contribution.
--
--   acc_adj += sum(acc_contribs)
--
-- I'm slightly worried about the asymptotics of this, since my
-- intuition of this is that the contributions might be rather sparse.
-- (Maybe completely zero?  If so it will be simplified away
-- entirely.)  Perhaps a better solution is to treat
-- accumulator-inputs in the primal code as we do free variables, and
-- create accumulators for them in the return sweep.
--
-- # Vectorised WithAcc
--
-- When WithAcc occurs in vectorised AD, the accumulator element types gain
-- extra leading "vectorised" dimensions corresponding to the enclosing vector
-- shape. For example, if the primal type inside a map of width @w@ is @acc(c,
-- [n], {f64})@, the adjoint type is @[w][n]f64@ -- but the internal accumulator
-- layout expects shape @[n][w]f64@ (the accumulator shape comes first, then the
-- vectorised dimensions, then element dimensions).
--
-- This means we must transpose accumulator adjoints when entering and
-- leaving the return-sweep WithAcc:
--
--  * On entry: transpose result adjoints from @[vec...][shape...]elem@ to
--    @[shape...][vec...]elem@ so they can serve as initial values for the
--    accumulators.
--
--  * On exit: transpose the produced arrays back from @[shape...][vec...]elem@
--    to @[vec...][shape...]elem@ to match the expected adjoint layout.
--
-- This is actually quite similar to how other SOACs must be handled.
--
-- Additionally, the accumulator parameter types in the lambda (and any
-- Acc-typed pattern elements or inner lambda parameters referring to the same
-- certs) must be updated to reflect the vectorised element types *before*
-- differentiation. This ensures that 'lookupAdj' on accumulator variables
-- inside the lambda produces adjoints with the correct vectorised type.
--
-- The UpdateAcc case is simpler under vectorisation: because the accumulator
-- adjoint already has the vectorised dimensions folded into its element type, a
-- plain index into the adjoint at the update indices directly yields the
-- correctly-shaped contribution.
--
-- # Consumption
--
-- A minor problem is that our usual way of handling consumption (Note
-- [Consumption]) is not viable, because accumulators are not
-- copyable.  Fortunately, while the accumulators that are consumed in
-- the forward sweep will also be present in the return sweep given
-- our current translation rules, they will be dead code.  As long as
-- we are careful to run dead code elimination after revVJP, we should
-- be good.
--
-- There is however one place where we must copy: the adjoint of the WithAcc
-- result serves double duty in the return sweep.  It is the initial value of
-- the accumulator (and thus consumed by the WithAcc), but it is *also* the
-- adjoint of the accumulator-typed result of the lambda body, and so may be
-- read inside the lambda - by the Map rule above, for example, which
-- replicates it.  These two uses are incompatible unless the WithAcc consumes
-- a copy.

-- Note [Array Adjoints of Match]
--
-- Some unusual, but sadly not completely contrived, contain Match
-- expressions that return multiple arrays, and there the arrays
-- returned by one branch have overlapping aliases with another
-- branch, although in different places. As an example consider this:
--
--   let (X,Y) = if c
--               then (A, B)
--               else (B, A)
--
-- Because our aliasing representation cannot express mutually
-- exclusive aliases, we will consider X and Y to be aliased to each
-- other. In practice, this means it is unlikely for X or Y to be
-- consumed, because it would also consume the other (although it's
-- possible for carefully written code).
--
-- When producing adjoints for this, it will be something like
--
--   let (X_adj,Y_adj) = if c
--                       then (A_adj, B_adj)
--                       else (B_adj, A_adj)
--
-- which completely reflects the primal code. However, while it is
-- unlikely that any consumption takes place for the original primal
-- variables, it is almost guaranteed that X_adj and Y_adj will be
-- consumed (that is the main way we use adjoints after all), and due
-- to the conservative aliasing, when one is consumed, so is the
-- other! To avoid this tragic fate, we are forced to copy any
-- array-typed adjoints returned by a Match. This can be quite costly.
-- However:
--
-- 1) Futhark has pretty OK copy removal, so maybe it can get rid of
--    these by using information not available to the AD pass.
--
-- 2) In many cases, arrays will have accumulator adjoints, which are
--    not subject to this problem.
--
-- Issue #2228 was caused by neglecting to do this.

import Control.Monad
import Control.Monad.Identity
import Data.List ((\\))
import Data.Maybe (isJust)
import Futhark.AD.Rev.Monad
import Futhark.AD.Shared (accAddLambda)
import Futhark.Builder
import Futhark.IR.SOACS
import Futhark.Tools
import Futhark.Transform.Rename
import Futhark.Util (chunks, takeLast)

-- | Transform updates on accumulators matching the given certificates into
-- updates that write provided zero values.
zeroOutUpdates :: [(VName, [SubExp])] -> Lambda SOACS -> Lambda SOACS
zeroOutUpdates :: [(VName, [SubExp])] -> Lambda SOACS -> Lambda SOACS
zeroOutUpdates [(VName, [SubExp])]
certs_to_zeroes Lambda SOACS
lam = Lambda SOACS
lam {lambdaBody = onBody $ lambdaBody lam}
  where
    onExp :: Exp SOACS -> Exp SOACS
onExp = Identity (Exp SOACS) -> Exp SOACS
forall a. Identity a -> a
runIdentity (Identity (Exp SOACS) -> Exp SOACS)
-> (Exp SOACS -> Identity (Exp SOACS)) -> Exp SOACS -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mapper SOACS SOACS Identity -> Exp SOACS -> Identity (Exp SOACS)
forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper SOACS SOACS Identity
mapper
      where
        mapper :: Mapper SOACS SOACS Identity
mapper =
          (Mapper SOACS SOACS m
forall rep (m :: * -> *). Monad m => Mapper rep rep m
forall {m :: * -> *}. Monad m => Mapper SOACS SOACS m
identityMapper :: (Monad m) => Mapper SOACS SOACS m)
            { mapOnOp = traverseSOACStms (\Scope SOACS
_ Stms SOACS
stms -> Stms SOACS -> Identity (Stms SOACS)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms SOACS -> Identity (Stms SOACS))
-> Stms SOACS -> Identity (Stms SOACS)
forall a b. (a -> b) -> a -> b
$ Stms SOACS -> Stms SOACS
onStms Stms SOACS
stms),
              mapOnBody = \Scope SOACS
_ Body SOACS
body -> Body SOACS -> Identity (Body SOACS)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Body SOACS -> Identity (Body SOACS))
-> Body SOACS -> Identity (Body SOACS)
forall a b. (a -> b) -> a -> b
$ Body SOACS -> Body SOACS
onBody Body SOACS
body
            }
    onStms :: Stms SOACS -> Stms SOACS
onStms = (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
onStm
    onStm :: Stm SOACS -> Stm SOACS
onStm (Let (Pat [PatElem (LetDec SOACS)
pe]) StmAux (ExpDec SOACS)
aux (BasicOp (UpdateAcc Safety
safety VName
acc [SubExp]
is [SubExp]
_)))
      | Acc VName
c Shape
_ [Type]
_ <- PatElem (LetDec SOACS) -> Type
forall dec. Typed dec => PatElem dec -> Type
patElemType PatElem (LetDec SOACS)
pe,
        Just [SubExp]
zero <- VName -> [(VName, [SubExp])] -> Maybe [SubExp]
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
c [(VName, [SubExp])]
certs_to_zeroes =
          Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let ([PatElem (LetDec SOACS)] -> Pat (LetDec SOACS)
forall dec. [PatElem dec] -> Pat dec
Pat [PatElem (LetDec SOACS)
pe]) StmAux (ExpDec SOACS)
aux (BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (Safety -> VName -> [SubExp] -> [SubExp] -> BasicOp
UpdateAcc Safety
safety VName
acc [SubExp]
is [SubExp]
zero))
    onStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux Exp SOACS
e) = Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$ Exp SOACS -> Exp SOACS
onExp Exp SOACS
e

    onBody :: Body SOACS -> Body SOACS
onBody Body SOACS
body = Body SOACS
body {bodyStms = onStms $ bodyStms body}

-- Update accumulator parameter types in the lambda to include vectorised
-- element types. Also updates all Acc-typed pattern elements and inner
-- lambda parameters that reference the same accumulator certs.
updateAccParamTypes :: Int -> Shape -> Lambda SOACS -> Lambda SOACS
updateAccParamTypes :: Int -> Shape -> Lambda SOACS -> Lambda SOACS
updateAccParamTypes Int
n_inputs Shape
adj_sh Lambda SOACS
lam
  | Shape
adj_sh Shape -> Shape -> Bool
forall a. Eq a => a -> a -> Bool
== Shape
forall a. Monoid a => a
mempty = Lambda SOACS
lam
  | Bool
otherwise =
      let ([Param Type]
cert_ps, [Param Type]
rest_ps) = Int -> [Param Type] -> ([Param Type], [Param Type])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n_inputs (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam)
          ([Param Type]
acc_ps, [Param Type]
other_ps) = Int -> [Param Type] -> ([Param Type], [Param Type])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n_inputs [Param Type]
rest_ps
          acc_ps' :: [Param Type]
acc_ps' = (Param Type -> Param Type) -> [Param Type] -> [Param Type]
forall a b. (a -> b) -> [a] -> [b]
map ([VName] -> Param Type -> Param Type
updateParam [VName]
cert_names) [Param Type]
acc_ps
          cert_names :: [VName]
cert_names = (Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName [Param Type]
cert_ps
          body' :: Body SOACS
body' = [VName] -> Body SOACS -> Body SOACS
updateBody [VName]
cert_names (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam)
          ret' :: [Type]
ret' = (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map ([VName] -> Type -> Type
updateAccType [VName]
cert_names) (Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
lam)
       in Lambda SOACS
lam
            { lambdaParams = cert_ps ++ acc_ps' ++ other_ps,
              lambdaReturnType = ret',
              lambdaBody = body'
            }
  where
    updateParam :: [VName] -> Param Type -> Param Type
    updateParam :: [VName] -> Param Type -> Param Type
updateParam [VName]
certs Param Type
p =
      Param Type
p {paramDec = updateAccType certs (paramDec p)}

    updateAccType :: [VName] -> Type -> Type
    updateAccType :: [VName] -> Type -> Type
updateAccType [VName]
certs (Acc VName
cert Shape
acc_shape [Type]
ts)
      | VName
cert VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
certs =
          VName -> Shape -> [Type] -> Type
forall shape o. VName -> Shape -> [Type] -> TypeBase shape o
Acc VName
cert Shape
acc_shape ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Type -> Shape -> Type
`arrayOfShape` Shape
adj_sh) [Type]
ts)
    updateAccType [VName]
_ Type
t = Type
t

    updateBody :: [VName] -> Body SOACS -> Body SOACS
    updateBody :: [VName] -> Body SOACS -> Body SOACS
updateBody [VName]
certs Body SOACS
body =
      Body SOACS
body {bodyStms = fmap (updateStm certs) (bodyStms body)}

    updateStm :: [VName] -> Stm SOACS -> Stm SOACS
    updateStm :: [VName] -> Stm SOACS -> Stm SOACS
updateStm [VName]
certs (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux Exp SOACS
e) =
      Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let ([VName] -> Pat Type -> Pat Type
updatePat [VName]
certs Pat Type
Pat (LetDec SOACS)
pat) StmAux (ExpDec SOACS)
aux ([VName] -> Exp SOACS -> Exp SOACS
updateExp [VName]
certs Exp SOACS
e)

    updatePat :: [VName] -> Pat Type -> Pat Type
    updatePat :: [VName] -> Pat Type -> Pat Type
updatePat [VName]
certs (Pat [PatElem Type]
pes) =
      [PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat ([PatElem Type] -> Pat Type) -> [PatElem Type] -> Pat Type
forall a b. (a -> b) -> a -> b
$ (PatElem Type -> PatElem Type) -> [PatElem Type] -> [PatElem Type]
forall a b. (a -> b) -> [a] -> [b]
map (\PatElem Type
pe -> PatElem Type
pe {patElemDec = updateAccType certs (patElemDec pe)}) [PatElem Type]
pes

    updateExp :: [VName] -> Exp SOACS -> Exp SOACS
    updateExp :: [VName] -> Exp SOACS -> Exp SOACS
updateExp [VName]
certs = Identity (Exp SOACS) -> Exp SOACS
forall a. Identity a -> a
runIdentity (Identity (Exp SOACS) -> Exp SOACS)
-> (Exp SOACS -> Identity (Exp SOACS)) -> Exp SOACS -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mapper SOACS SOACS Identity -> Exp SOACS -> Identity (Exp SOACS)
forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper SOACS SOACS Identity
mapper
      where
        mapper :: Mapper SOACS SOACS Identity
mapper =
          (Mapper SOACS SOACS m
forall rep (m :: * -> *). Monad m => Mapper rep rep m
forall {m :: * -> *}. Monad m => Mapper SOACS SOACS m
identityMapper :: (Monad m) => Mapper SOACS SOACS m)
            { mapOnBody = \Scope SOACS
_ Body SOACS
b -> Body SOACS -> Identity (Body SOACS)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Body SOACS -> Identity (Body SOACS))
-> Body SOACS -> Identity (Body SOACS)
forall a b. (a -> b) -> a -> b
$ [VName] -> Body SOACS -> Body SOACS
updateBody [VName]
certs Body SOACS
b,
              mapOnOp = pure . updateSOAC certs
            }

    updateSOAC :: [VName] -> SOAC SOACS -> SOAC SOACS
    updateSOAC :: [VName] -> SOAC SOACS -> SOAC SOACS
updateSOAC [VName]
certs = Identity (SOAC SOACS) -> SOAC SOACS
forall a. Identity a -> a
runIdentity (Identity (SOAC SOACS) -> SOAC SOACS)
-> (SOAC SOACS -> Identity (SOAC SOACS))
-> SOAC SOACS
-> SOAC SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOACMapper SOACS SOACS Identity
-> SOAC SOACS -> Identity (SOAC SOACS)
forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM SOACMapper SOACS SOACS Identity
mapper
      where
        mapper :: SOACMapper SOACS SOACS Identity
mapper =
          SOACMapper SOACS SOACS Identity
forall rep (m :: * -> *). Monad m => SOACMapper rep rep m
identitySOACMapper
            { mapOnSOACLambda = pure . updateLambda certs
            }

    updateLambda :: [VName] -> Lambda SOACS -> Lambda SOACS
    updateLambda :: [VName] -> Lambda SOACS -> Lambda SOACS
updateLambda [VName]
certs Lambda SOACS
l =
      Lambda SOACS
l
        { lambdaParams = map (updateParam certs) (lambdaParams l),
          lambdaReturnType = map (updateAccType certs) (lambdaReturnType l),
          lambdaBody = updateBody certs (lambdaBody l)
        }

diffWithAcc ::
  VjpOps ->
  Pat Type ->
  StmAux () ->
  [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))] ->
  Lambda SOACS ->
  ADM () ->
  ADM ()
diffWithAcc :: VjpOps
-> Pat Type
-> StmAux ()
-> [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
-> Lambda SOACS
-> ADM ()
-> ADM ()
diffWithAcc VjpOps
ops Pat Type
pat StmAux ()
aux [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs Lambda SOACS
lam 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
$ [WithAccInput (Rep ADM)] -> Lambda (Rep ADM) -> Exp (Rep ADM)
forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc [WithAccInput (Rep ADM)]
[(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs Lambda (Rep ADM)
Lambda SOACS
lam
  ADM ()
m
  ADM () -> ADM ()
forall a. ADM a -> ADM a
returnSweepCode (ADM () -> ADM ()) -> ADM () -> ADM ()
forall a b. (a -> b) -> a -> b
$ do
    adj_shape <- ADM Shape
askShape
    adjs <- mapM lookupAdj $ patNames pat
    -- Transpose the accumulator result adjoints from [vec...][shape...]elem
    -- to [shape...][vec...]elem, matching the internal accumulator layout.
    adjs' <- transposeAdjs adj_shape adjs
    lam' <- renameLambda lam
    -- Update the lambda's accumulator parameter types to reflect vectorised
    -- element types BEFORE differentiation, so that lookupAdj on Acc variables
    -- inside the lambda gives the correct vectorised adjoint type.
    let lam'_vec = Int -> Shape -> Lambda SOACS -> Lambda SOACS
updateAccParamTypes Int
n_inputs Shape
adj_shape Lambda SOACS
lam'
    free_vars <- filterM isActive $ namesToList $ freeIn lam'_vec
    free_accs <- filterM (fmap isAcc . lookupType) free_vars
    let free_vars' = [VName]
free_vars [VName] -> [VName] -> [VName]
forall a. Eq a => [a] -> [a] -> [a]
\\ [VName]
free_accs
        op_certs =
          ((Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))
 -> VName)
-> [(Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))]
-> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (Param Type -> VName
forall dec. Param dec -> VName
paramName (Param Type -> VName)
-> ((Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))
    -> Param Type)
-> (Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))
-> Param Type
forall a b. (a, b) -> a
fst)
            ([(Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))]
 -> [VName])
-> ([(Param Type,
      (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))]
    -> [(Param Type,
         (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))])
-> [(Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))]
-> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))
 -> Bool)
-> [(Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))]
-> [(Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Shape, [VName], Maybe (Lambda SOACS, [SubExp])) -> Bool
forall {a} {b} {a}. (a, b, Maybe a) -> Bool
hasOperator ((Shape, [VName], Maybe (Lambda SOACS, [SubExp])) -> Bool)
-> ((Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))
    -> (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))
-> (Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))
-> (Shape, [VName], Maybe (Lambda SOACS, [SubExp]))
forall a b. (a, b) -> b
snd)
            ([(Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))]
 -> [VName])
-> [(Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))]
-> [VName]
forall a b. (a -> b) -> a -> b
$ [Param Type]
-> [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
-> [(Param Type, (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))]
forall a b. [a] -> [b] -> [(a, b)]
zip (Int -> [Param Type] -> [Param Type]
forall a. Int -> [a] -> [a]
take Int
n_inputs (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam'_vec)) [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs
    lam'' <- withOperatorAccs op_certs $ diffLambda' adjs' free_vars' lam'_vec
    (inputs_zeroes, inputs') <-
      unzip <$> zipWithM (renameInputLambda adj_shape) (chunks lengths adjs) inputs
    let certs = (Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName ([Param Type] -> [VName]) -> [Param Type] -> [VName]
forall a b. (a -> b) -> a -> b
$ Int -> [Param Type] -> [Param Type]
forall a. Int -> [a] -> [a]
take Int
n_inputs ([Param Type] -> [Param Type]) -> [Param Type] -> [Param Type]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam''
    raw_adjs <-
      letTupExp "with_acc_contrib" . WithAcc inputs' $
        zeroOutUpdates (zip certs inputs_zeroes) lam''
    -- The accumulator results have shape [shape...][vec...]elem. Transpose
    -- back to [vec...][shape...]elem for the adjoint.
    let n_arrs = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
lengths
        (arr_adjs, free_adjs) = splitAt n_arrs raw_adjs
    arr_adjs' <- zipWithM (transposeAccResult adj_shape) (map (\(Shape
s, [VName]
_, Maybe (Lambda SOACS, [SubExp])
_) -> Shape
s) inputs) arr_adjs
    zipWithM_ insAdj arrs arr_adjs'
    zipWithM_ insAdj free_vars' free_adjs
  where
    n_inputs :: Int
n_inputs = [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs
    lengths :: [Int]
lengths = ((Shape, [VName], Maybe (Lambda SOACS, [SubExp])) -> Int)
-> [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (\(Shape
_, [VName]
as, Maybe (Lambda SOACS, [SubExp])
_) -> [VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
as) [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs
    arrs :: [VName]
arrs = ((Shape, [VName], Maybe (Lambda SOACS, [SubExp])) -> [VName])
-> [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(Shape
_, [VName]
as, Maybe (Lambda SOACS, [SubExp])
_) -> [VName]
as) [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs
    hasOperator :: (a, b, Maybe a) -> Bool
hasOperator (a
_, b
_, Maybe a
op) = Maybe a -> Bool
forall a. Maybe a -> Bool
isJust Maybe a
op

    -- Transpose the accumulator-related adjoints from [vec...][shape...]elem
    -- to [shape...][vec...]elem. Non-accumulator adjs are left unchanged.
    transposeAdjs :: Shape -> [Adj] -> ADM [Adj]
    transposeAdjs :: Shape -> [Adj] -> ADM [Adj]
transposeAdjs Shape
adj_sh [Adj]
adjs
      | Shape
adj_sh Shape -> Shape -> Bool
forall a. Eq a => a -> a -> Bool
== Shape
forall a. Monoid a => a
mempty = [Adj] -> ADM [Adj]
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Adj]
adjs
      | Bool
otherwise = do
          let n_arrs :: Int
n_arrs = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
lengths
              ([Adj]
acc_adjs, [Adj]
other_adjs) = Int -> [Adj] -> ([Adj], [Adj])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n_arrs [Adj]
adjs
          acc_adjs' <- (Adj -> ADM Adj) -> [Adj] -> 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 Adj -> ADM Adj
transposeAdj [Adj]
acc_adjs
          pure $ acc_adjs' ++ other_adjs

    transposeAdj :: Adj -> ADM Adj
    transposeAdj :: Adj -> ADM Adj
transposeAdj Adj
adj = do
      v <- Adj -> ADM VName
adjVal Adj
adj
      v' <- vecToInner v
      pure $ AdjVal $ Var v'

    -- Transpose [shape...][vec...][elem...] to [vec...][shape...][elem...]
    transposeAccResult :: Shape -> Shape -> VName -> ADM VName
    transposeAccResult :: Shape -> Shape -> VName -> ADM VName
transposeAccResult Shape
adj_sh Shape
shape VName
v
      | Shape
adj_sh Shape -> Shape -> Bool
forall a. Eq a => a -> a -> Bool
== Shape
forall a. Monoid a => a
mempty = VName -> ADM VName
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
      | Bool
otherwise = do
          v_t <- VName -> ADM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
          let r = Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
adj_sh
              s = Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
shape
              total = Type -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
arrayRank Type
v_t
              perm = [Int
s .. Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
0 .. Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r .. Int
total Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
          letExp (baseName v <> "_tr") $ BasicOp $ Rearrange v perm

    renameInputLambda :: Shape
-> t Adj
-> (a, [VName], Maybe a)
-> ADM ([SubExp], (a, t VName, Maybe (Lambda SOACS, [SubExp])))
renameInputLambda Shape
adj_sh t Adj
as_adj (a
shape, [VName]
as, Maybe a
op) = do
      -- Compute element types with vectorised dimensions included.
      orig_nes_ts <- (VName -> ADM Type) -> [VName] -> ADM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((Type -> Type) -> ADM Type -> ADM Type
forall a b. (a -> b) -> ADM a -> ADM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> Type -> Type
forall d m.
Int -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
stripArray (a -> Int
forall a. ArrayShape a => a -> Int
shapeRank a
shape)) (ADM Type -> ADM Type) -> (VName -> ADM Type) -> VName -> ADM Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> ADM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType) [VName]
as
      let vec_nes_ts = (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Type -> Shape -> Type
`arrayOfShape` Shape
adj_sh) [Type]
orig_nes_ts
      zeroes <- mapM (zeroArray mempty) vec_nes_ts
      -- The result adjoint is consumed by the WithAcc, but is also the adjoint
      -- of the accumulator inside the lambda (see Note [Adjoints of
      -- accumulators]), so it must remain readable there. Hence the copy.
      as' <- mapM (letExp "acc_adj_init" . BasicOp . Replicate mempty . Var <=< adjVal) as_adj
      -- Transpose adjoints from [vec...][shape...]elem to [shape...][vec...]elem
      -- so they match the accumulator layout.
      as'' <- mapM vecToInner as'
      -- 'zeroOutUpdates' makes the primal updates write zeroes. For a
      -- scatter-like accumulator that overwrites the cell, this is what we
      -- want. An accumulator with a combining operator does not overwrite, so
      -- we give this one addition as its operator, making those writes no-ops.
      op' <- case op of
        Maybe a
Nothing -> Maybe (Lambda SOACS, [SubExp])
-> ADM (Maybe (Lambda SOACS, [SubExp]))
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Lambda SOACS, [SubExp])
forall a. Maybe a
Nothing
        -- Under vectorisation the element types gain the vector dimensions, so
        -- the original operator no longer fits. We assume it is addition
        -- anyway, so just build that.
        Just a
_ -> do
          add_lam <- Int -> [Type] -> ADM (Lambda SOACS)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ SOACS) =>
Int -> [Type] -> m (Lambda SOACS)
accAddLambda (a -> Int
forall a. ArrayShape a => a -> Int
shapeRank a
shape) [Type]
vec_nes_ts
          nes <- mapM (letSubExp "acc_adj_zero" . zeroExp) vec_nes_ts
          pure $ Just (add_lam, nes)
      pure (map Var zeroes, (shape, as'', op'))

    diffLambda' :: [Adj] -> [VName] -> GLambda SOACS t -> ADM (Lambda (Rep ADM))
diffLambda' [Adj]
res_adjs [VName]
get_adjs_for (Lambda [LParam SOACS]
params [t]
_ Body SOACS
body) =
      [LParam (Rep ADM)] -> ADM Result -> ADM (Lambda (Rep ADM))
forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
mkLambda [LParam (Rep ADM)]
[LParam SOACS]
params (ADM Result -> ADM (Lambda (Rep ADM)))
-> ADM Result -> ADM (Lambda (Rep ADM))
forall a b. (a -> b) -> a -> b
$ do
        res <- Body (Rep ADM) -> ADM Result
Body SOACS -> ADM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind (Body SOACS -> ADM Result) -> ADM (Body SOACS) -> ADM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VjpOps -> [Adj] -> [VName] -> Body SOACS -> ADM (Body SOACS)
vjpBody VjpOps
ops [Adj]
res_adjs [VName]
get_adjs_for Body SOACS
body
        pure $ take n_inputs res <> takeLast (length get_adjs_for) res

diffUpdateAcc ::
  Pat Type ->
  StmAux () ->
  Safety ->
  VName ->
  [SubExp] ->
  [SubExp] ->
  ADM () ->
  ADM ()
diffUpdateAcc :: 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 = do
  -- By the type rules for UpdateAcc, the pattern must be a singleton.
  let Pat ~[PatElem Type
pe] = Pat Type
pat
  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 -> Exp (Rep ADM)) -> BasicOp -> Exp (Rep ADM)
forall a b. (a -> b) -> a -> b
$ Safety -> VName -> [SubExp] -> [SubExp] -> BasicOp
UpdateAcc Safety
safety VName
acc [SubExp]
is [SubExp]
vs
  ADM ()
m
  adj <- VName -> ADM VName
lookupAdjVal (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
  returnSweepCode $ do
    adj_t <- lookupType adj
    acc_t <- lookupType acc
    -- An accumulator with a combining operator does not overwrite, so the
    -- incoming value of the updated cell retains its full sensitivity.
    overwrites <- case acc_t of
      Acc VName
cert Shape
_ [Type]
_ -> Bool -> Bool
not (Bool -> Bool) -> ADM Bool -> ADM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> ADM Bool
isOperatorAcc VName
cert
      Type
_ -> Bool -> ADM Bool
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
    let elem_t = Int -> Type -> Type
forall d m.
Int -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
stripArray ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
is) Type
adj_t
        slice = Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice Type
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
forall d. d -> DimIndex d
DimFix [SubExp]
is
        -- The value adjoint is the corresponding cell of the accumulator
        -- adjoint.
        index_adj = ADM (Exp (Rep ADM)) -> ADM (Exp (Rep ADM))
maybe_copy (ADM (Exp (Rep ADM)) -> ADM (Exp (Rep ADM)))
-> ADM (Exp (Rep ADM)) -> ADM (Exp (Rep ADM))
forall a b. (a -> b) -> a -> b
$ Exp (Rep ADM) -> ADM (Exp (Rep ADM))
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep ADM) -> ADM (Exp (Rep ADM)))
-> Exp (Rep ADM) -> ADM (Exp (Rep ADM))
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
$ VName -> Slice SubExp -> BasicOp
Index VName
adj Slice SubExp
slice
          where
            -- We have to copy a slice because we are updating 'adj' as well -
            -- even though in many cases that update is likely dead code... Not
            -- great.
            maybe_copy :: ADM (Exp (Rep ADM)) -> ADM (Exp (Rep ADM))
maybe_copy
              | [SubExp] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([SubExp] -> Bool) -> [SubExp] -> Bool
forall a b. (a -> b) -> a -> b
$ Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice = ADM (Exp (Rep ADM)) -> ADM (Exp (Rep ADM))
forall a. a -> a
id
              | Bool
otherwise = ADM (Exp (Rep ADM)) -> ADM (Exp (Rep ADM))
forall (m :: * -> *).
MonadBuilder m =>
m (Exp (Rep m)) -> m (Exp (Rep m))
eCopy
        -- For a scatter-like accumulator, the input accumulator adjoint is the
        -- result adjoint with the updated cell zeroed out: a cell that is
        -- subsequently overwritten does not contribute to the result, and so
        -- has zero sensitivity.
        zeroed
          | Bool
overwrites = do
              z <- Name -> Exp (Rep ADM) -> ADM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"acc_adj_zero" (Exp (Rep ADM) -> ADM SubExp) -> Exp (Rep ADM) -> ADM SubExp
forall a b. (a -> b) -> a -> b
$ Type -> Exp (Rep ADM)
forall rep. Type -> Exp rep
zeroExp Type
elem_t
              pure $ BasicOp $ Update Unsafe adj slice z
          | Bool
otherwise = Exp rep -> ADM (Exp rep)
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp rep -> ADM (Exp rep)) -> Exp rep -> ADM (Exp rep)
forall a b. (a -> b) -> a -> b
$ 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
$ VName -> SubExp
Var VName
adj
    (adj_i, acc_adj) <- case safety of
      Safety
Unsafe ->
        (,)
          (VName -> VName -> (VName, VName))
-> ADM VName -> ADM (VName -> (VName, VName))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Name -> Exp (Rep ADM) -> ADM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"updateacc_val_adj" (Exp SOACS -> ADM VName) -> ADM (Exp SOACS) -> ADM VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ADM (Exp (Rep ADM))
ADM (Exp SOACS)
index_adj)
          ADM (VName -> (VName, VName)) -> ADM VName -> ADM (VName, VName)
forall a b. ADM (a -> b) -> ADM a -> ADM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Name -> Exp (Rep ADM) -> ADM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"acc_adj" (Exp SOACS -> ADM VName) -> ADM (Exp SOACS) -> ADM VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ADM (Exp SOACS)
forall {rep}. ADM (Exp rep)
zeroed)
      Safety
Safe -> do
        -- The primal UpdateAcc may be out-of-bounds, in which case indexing the
        -- adjoint is dangerous and the input accumulator adjoint is unchanged.
        ~[adj_i, acc_adj] <-
          Name -> Exp (Rep ADM) -> ADM [VName]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [VName]
letTupExp Name
"updateacc_adj"
            (Exp SOACS -> ADM [VName]) -> ADM (Exp SOACS) -> ADM [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ADM (Exp (Rep ADM))
-> ADM (Body (Rep ADM))
-> ADM (Body (Rep ADM))
-> ADM (Exp (Rep ADM))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
              (Shape -> [ADM (Exp (Rep ADM))] -> ADM (Exp (Rep ADM))
forall (m :: * -> *).
MonadBuilder m =>
Shape -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eShapeInBounds (Type -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
adj_t) ((SubExp -> ADM (Exp SOACS)) -> [SubExp] -> [ADM (Exp SOACS)]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> ADM (Exp (Rep ADM))
SubExp -> ADM (Exp SOACS)
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp [SubExp]
is))
              ([ADM (Exp (Rep ADM))] -> ADM (Body (Rep ADM))
forall (m :: * -> *).
MonadBuilder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody [ADM (Exp (Rep ADM))
index_adj, ADM (Exp (Rep ADM))
ADM (Exp SOACS)
forall {rep}. ADM (Exp rep)
zeroed])
              ([ADM (Exp (Rep ADM))] -> ADM (Body (Rep ADM))
forall (m :: * -> *).
MonadBuilder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody [Exp (Rep ADM) -> ADM (Exp (Rep ADM))
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep ADM) -> ADM (Exp (Rep ADM)))
-> Exp (Rep ADM) -> ADM (Exp (Rep ADM))
forall a b. (a -> b) -> a -> b
$ Type -> Exp (Rep ADM)
forall rep. Type -> Exp rep
zeroExp Type
elem_t, Exp (Rep ADM) -> ADM (Exp (Rep ADM))
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep ADM) -> ADM (Exp (Rep ADM)))
-> Exp (Rep ADM) -> ADM (Exp (Rep ADM))
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 -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
adj])
        pure (adj_i, acc_adj)
    -- XXX: this is only OK because we assume accumulators are currently
    -- singleton.
    updateSubExpAdj (head vs) adj_i
    insAdj acc acc_adj