{-# LANGUAGE TypeFamilies #-}

-- | VJP transformation for 'Map' and 'FlatMap'. This is a pretty complicated
-- case due to the possibility of free variables. The two are handled together
-- because the return sweep of a 'FlatMap' is itself a 'Map', and so requires
-- all the same machinery.
module Futhark.AD.Rev.Map (vjpMap, vjpFlatMap) where

import Control.Monad
import Data.Bifunctor (first, second)
import Data.Either (rights)
import Data.Maybe (catMaybes)
import Futhark.AD.Rev.Monad
import Futhark.AD.Shared (accAddLambda, asVName, vecPerm)
import Futhark.Analysis.PrimExp.Convert
import Futhark.Builder
import Futhark.IR.SOACS
import Futhark.Tools hiding (withAcc)
import Futhark.Transform.Rename
import Futhark.Util (splitAt3)

-- | A classification of a free variable based on its adjoint.  The
-- 'VName' stored is *not* the adjoint, but the primal variable.
data AdjVar
  = -- | Adjoint is already an accumulator.
    FreeAcc VName
  | -- | Currently has no adjoint, but should be given one, and is an
    -- array with this shape and element type.
    FreeArr VName Shape PrimType
  | -- | Does not need an accumulator adjoint (might still be an array).
    FreeNonAcc VName

classifyAdjVars :: [VName] -> ADM [AdjVar]
classifyAdjVars :: [VName] -> ADM [AdjVar]
classifyAdjVars = (VName -> ADM AdjVar) -> [VName] -> ADM [AdjVar]
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 AdjVar
f
  where
    f :: VName -> ADM AdjVar
f VName
v = do
      v_adj <- VName -> ADM VName
lookupAdjVal VName
v
      v_adj_t <- lookupType v_adj
      case v_adj_t of
        Array PrimType
pt Shape
shape NoMode
_ ->
          AdjVar -> ADM AdjVar
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AdjVar -> ADM AdjVar) -> AdjVar -> ADM AdjVar
forall a b. (a -> b) -> a -> b
$ VName -> Shape -> PrimType -> AdjVar
FreeArr VName
v Shape
shape PrimType
pt
        Acc {} ->
          AdjVar -> ADM AdjVar
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AdjVar -> ADM AdjVar) -> AdjVar -> ADM AdjVar
forall a b. (a -> b) -> a -> b
$ VName -> AdjVar
FreeAcc VName
v
        Type
_ ->
          AdjVar -> ADM AdjVar
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AdjVar -> ADM AdjVar) -> AdjVar -> ADM AdjVar
forall a b. (a -> b) -> a -> b
$ VName -> AdjVar
FreeNonAcc VName
v

partitionAdjVars :: [AdjVar] -> ([(VName, (Shape, PrimType))], [VName], [VName])
partitionAdjVars :: [AdjVar] -> ([(VName, (Shape, PrimType))], [VName], [VName])
partitionAdjVars [] = ([], [], [])
partitionAdjVars (AdjVar
fv : [AdjVar]
fvs) =
  case AdjVar
fv of
    FreeArr VName
v Shape
shape PrimType
t -> ((VName
v, (Shape
shape, PrimType
t)) (VName, (Shape, PrimType))
-> [(VName, (Shape, PrimType))] -> [(VName, (Shape, PrimType))]
forall a. a -> [a] -> [a]
: [(VName, (Shape, PrimType))]
xs, [VName]
ys, [VName]
zs)
    FreeAcc VName
v -> ([(VName, (Shape, PrimType))]
xs, VName
v VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: [VName]
ys, [VName]
zs)
    FreeNonAcc VName
v -> ([(VName, (Shape, PrimType))]
xs, [VName]
ys, VName
v VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: [VName]
zs)
  where
    ([(VName, (Shape, PrimType))]
xs, [VName]
ys, [VName]
zs) = [AdjVar] -> ([(VName, (Shape, PrimType))], [VName], [VName])
partitionAdjVars [AdjVar]
fvs

buildRenamedBody ::
  (MonadBuilder m) =>
  m (Result, a) ->
  m (Body (Rep m), a)
buildRenamedBody :: forall (m :: * -> *) a.
MonadBuilder m =>
m (Result, a) -> m (Body (Rep m), a)
buildRenamedBody m (Result, a)
m = do
  (body, x) <- m (Result, a) -> m (GBody (Rep m) SubExpRes, a)
forall (m :: * -> *) res a.
(MonadBuilder m, IsResult res) =>
m ([res], a) -> m (GBody (Rep m) res, a)
buildBody m (Result, a)
m
  body' <- renameBody body
  pure (body', x)

withAcc ::
  [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))] ->
  ([VName] -> ADM Result) ->
  ADM [VName]
withAcc :: [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
-> ([VName] -> ADM Result) -> ADM [VName]
withAcc [] [VName] -> ADM Result
m =
  (SubExpRes -> ADM VName) -> Result -> 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 (Name -> Exp (Rep ADM) -> ADM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"withacc_res" (Exp SOACS -> ADM VName)
-> (SubExpRes -> Exp SOACS) -> SubExpRes -> ADM VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS)
-> (SubExpRes -> BasicOp) -> SubExpRes -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp (SubExp -> BasicOp)
-> (SubExpRes -> SubExp) -> SubExpRes -> BasicOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExpRes -> SubExp
resSubExp) (Result -> ADM [VName]) -> ADM Result -> ADM [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [VName] -> ADM Result
m []
withAcc [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs [VName] -> ADM Result
m = do
  (cert_params, acc_params) <- ([(Param Type, Param Type)] -> ([Param Type], [Param Type]))
-> ADM [(Param Type, Param Type)]
-> ADM ([Param Type], [Param Type])
forall a b. (a -> b) -> ADM a -> ADM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(Param Type, Param Type)] -> ([Param Type], [Param Type])
forall a b. [(a, b)] -> ([a], [b])
unzip (ADM [(Param Type, Param Type)]
 -> ADM ([Param Type], [Param Type]))
-> ADM [(Param Type, Param Type)]
-> ADM ([Param Type], [Param Type])
forall a b. (a -> b) -> a -> b
$
    [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
-> ((Shape, [VName], Maybe (Lambda SOACS, [SubExp]))
    -> ADM (Param Type, Param Type))
-> ADM [(Param Type, Param Type)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs (((Shape, [VName], Maybe (Lambda SOACS, [SubExp]))
  -> ADM (Param Type, Param Type))
 -> ADM [(Param Type, Param Type)])
-> ((Shape, [VName], Maybe (Lambda SOACS, [SubExp]))
    -> ADM (Param Type, Param Type))
-> ADM [(Param Type, Param Type)]
forall a b. (a -> b) -> a -> b
$ \(Shape
shape, [VName]
arrs, Maybe (Lambda SOACS, [SubExp])
_) -> do
      cert_param <- Name -> Type -> ADM (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"acc_cert_p" (Type -> ADM (Param Type)) -> Type -> ADM (Param Type)
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Unit
      ts <- mapM (fmap (stripArray (shapeRank shape)) . lookupType) arrs
      acc_param <- newParam "acc_p" $ Acc (paramName cert_param) shape ts
      pure (cert_param, acc_param)
  acc_lam <-
    subAD $ mkLambda (cert_params ++ acc_params) $ m $ map paramName acc_params
  letTupExp "withhacc_res" $ WithAcc inputs acc_lam

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

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

withAccInput ::
  (VName, (Shape, PrimType)) ->
  ADM (Shape, [VName], Maybe (Lambda SOACS, [SubExp]))
withAccInput :: (VName, (Shape, PrimType))
-> ADM (Shape, [VName], Maybe (Lambda SOACS, [SubExp]))
withAccInput (VName
v, (Shape
shape, PrimType
pt)) = do
  v_adj <- VName -> ADM VName
lookupAdjVal VName
v
  add_lam <- accAddLambda (shapeRank shape) [Prim pt]
  zero <- letSubExp "zero" $ zeroExp $ Prim pt
  pure (shape, [v_adj], Just (add_lam, [zero]))

-- | Run an action in a context where the array-typed adjoints of the given
-- free variables have been turned into accumulators, so that the contributions
-- from each iteration of the SOAC are summed.  The action is passed the free
-- variables that were given accumulator adjoints, and those that were not.  The
-- 'VName' list is the input arrays of the SOAC.
accAdjoints :: [VName] -> [VName] -> ([VName] -> Names -> ADM ()) -> ADM ()
accAdjoints :: [VName] -> [VName] -> ([VName] -> Names -> ADM ()) -> ADM ()
accAdjoints [VName]
as [VName]
free [VName] -> Names -> ADM ()
m = do
  (arr_free, acc_free, nonacc_free) <-
    [AdjVar] -> ([(VName, (Shape, PrimType))], [VName], [VName])
partitionAdjVars ([AdjVar] -> ([(VName, (Shape, PrimType))], [VName], [VName]))
-> ADM [AdjVar]
-> ADM ([(VName, (Shape, PrimType))], [VName], [VName])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [VName] -> ADM [AdjVar]
classifyAdjVars [VName]
free
  arr_free' <- mapM withAccInput arr_free
  -- We only consider those input arrays that are also not free in
  -- the lambda.
  let as_nonfree = (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]
free) [VName]
as
  (arr_adjs, acc_adjs, rest_adjs) <-
    fmap (splitAt3 (length arr_free) (length acc_free)) . withAcc arr_free' $ \[VName]
accs -> 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, (Shape, PrimType)) -> VName)
-> [(VName, (Shape, PrimType))] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, (Shape, PrimType)) -> VName
forall a b. (a, b) -> a
fst [(VName, (Shape, PrimType))]
arr_free) [VName]
accs
      () <- [VName] -> Names -> ADM ()
m ([VName]
acc_free [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ ((VName, (Shape, PrimType)) -> VName)
-> [(VName, (Shape, PrimType))] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, (Shape, PrimType)) -> VName
forall a b. (a, b) -> a
fst [(VName, (Shape, PrimType))]
arr_free) ([VName] -> Names
namesFromList [VName]
nonacc_free)
      acc_free_adj <- mapM lookupAdjVal acc_free
      arr_free_adj <- mapM (lookupAdjVal . fst) arr_free
      nonacc_free_adj <- mapM lookupAdjVal nonacc_free
      as_nonfree_adj <- mapM lookupAdjVal as_nonfree
      pure $ varsRes $ arr_free_adj <> acc_free_adj <> nonacc_free_adj <> as_nonfree_adj
  zipWithM_ insAdj acc_free acc_adjs
  zipWithM_ insAdj (map fst arr_free) arr_adjs
  let (nonacc_adjs, as_nonfree_adjs) = splitAt (length nonacc_free) rest_adjs
  zipWithM_ insAdj nonacc_free nonacc_adjs
  zipWithM_ insAdj as_nonfree as_nonfree_adjs

-- | Add the per-iteration contributions to the adjoint of a free variable.  If
-- the adjoint is an accumulator, the summation has already taken place.
freeContrib :: SubExp -> VName -> VName -> ADM ()
freeContrib :: SubExp -> VName -> VName -> ADM ()
freeContrib SubExp
w VName
v VName
contribs = do
  contribs_t <- VName -> ADM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
contribs
  case rowType contribs_t of
    Acc {} -> 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 ()
insAdj VName
v VName
contribs
    Type
t -> do
      lam <- Type -> ADM (Lambda (Rep ADM))
forall (m :: * -> *).
(OpC (Rep m) ~ SOAC, MonadBuilder m, Buildable (Rep m)) =>
Type -> m (Lambda (Rep m))
addLambda Type
t
      zero <- letSubExp "zero" $ zeroExp t
      reduce <- reduceSOAC [Reduce Commutative lam [zero]]
      contrib_sum <-
        letExp (baseName v <> "_contrib_sum") . Op $
          Screma w [contribs] reduce
      void $ updateAdj v contrib_sum

-- | Turn the 'ExtLambda' of a 'FlatMap' into an ordinary 'Lambda', by dropping
-- the size result and coercing each nonuniform result to the given size, which
-- must be dynamically equal to the size the lambda computes.  See Note
-- [Adjoints of FlatMap].
flatMapPlainLambda :: SubExp -> ExtLambda SOACS -> ADM (Lambda SOACS)
flatMapPlainLambda :: SubExp -> ExtLambda SOACS -> ADM (Lambda SOACS)
flatMapPlainLambda SubExp
n (Lambda [LParam SOACS]
params [ExtType]
rettype 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
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body (Rep ADM)
Body SOACS
body
    forM (zip (drop 1 rettype) (drop 1 res)) $ \(ExtType
t, res_i :: SubExpRes
res_i@(SubExpRes Certs
cs SubExp
se)) ->
      if ExtType -> Bool
flatMapNonuniform ExtType
t
        then do
          v <- SubExp -> ADM VName
forall (m :: * -> *). MonadBuilder m => SubExp -> m VName
asVName SubExp
se
          v_t <- lookupType v
          fmap varRes . certifying cs . letExp "flatmap_res_coerce" $
            shapeCoerce (arrayDims (v_t `setOuterSize` n)) v
        else SubExpRes -> ADM SubExpRes
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExpRes
res_i

-- | Construct the Map that constitutes the return sweep of a Map-like SOAC.
-- Contributions to the free variables of the lambda are handled here; the
-- contributions to the input arrays are passed to the continuation.
mapReturnSweep ::
  (FreeIn t) =>
  VjpOps ->
  StmAux () ->
  -- | Width of the SOAC and its input arrays.
  (SubExp, [VName]) ->
  -- | Lambda of the SOAC, which must already have been renamed.  Only its
  -- parameters and free variables are used.
  GLambda SOACS t ->
  -- | Additional arrays to map across, and the parameters receiving their
  -- elements.
  [(VName, LParam SOACS)] ->
  -- | Produce the adjoints of the results of the lambda to differentiate, along
  -- with that lambda - which need not be the lambda of the SOAC, but must have
  -- the same parameters.  Run inside the return sweep lambda.
  ADM ([Adj], Lambda SOACS) ->
  -- | Given the contribution to each input array.
  ([VName] -> ADM ()) ->
  ADM ()
mapReturnSweep :: forall t.
FreeIn t =>
VjpOps
-> StmAux ()
-> (SubExp, [VName])
-> GLambda SOACS t
-> [(VName, LParam SOACS)]
-> ADM ([Adj], Lambda SOACS)
-> ([VName] -> ADM ())
-> ADM ()
mapReturnSweep VjpOps
ops StmAux ()
aux (SubExp
w, [VName]
as) GLambda SOACS t
lam [(VName, LParam SOACS)]
extra ADM ([Adj], Lambda SOACS)
mkAdjs [VName] -> ADM ()
onContribs = do
  free <- (VName -> ADM Bool) -> [VName] -> ADM [VName]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM VName -> ADM Bool
isActive ([VName] -> ADM [VName]) -> [VName] -> ADM [VName]
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ GLambda SOACS t -> Names
forall a. FreeIn a => a -> Names
freeIn GLambda SOACS t
lam
  accAdjoints as free $ \[VName]
free_with_adjs Names
free_without_adjs -> do
    free_adjs <- (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]
free_with_adjs
    free_adjs_ts <- mapM lookupType free_adjs
    free_adjs_params <- mapM (newParam "free_adj_p") free_adjs_ts
    let (extra_arrs, extra_params) = unzip extra
        adjs_for = (Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName (GLambda SOACS t -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams GLambda SOACS t
lam) [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
free
    lam_rev <-
      mkLambda (lambdaParams lam ++ extra_params ++ free_adjs_params)
        . subAD
        . noAdjsFor free_without_adjs
        $ do
          zipWithM_ insAdj free_with_adjs $ map paramName free_adjs_params
          (res_adjs, lam') <- mkAdjs
          bodyBind . lambdaBody =<< vjpLambda ops res_adjs adjs_for lam'

    (param_contribs, free_contribs) <-
      fmap (splitAt (length (lambdaParams lam))) $
        auxing aux
          . letTupExp "map_adjs"
          . Op
          . Screma w (as ++ extra_arrs ++ free_adjs)
          =<< mapSOAC lam_rev

    -- Crucial that we handle the free contribs first in case 'free'
    -- and 'as' intersect.
    zipWithM_ (freeContrib w) free free_contribs
    onContribs param_contribs

-- | Perform VJP on a Map.  The 'Adj' list is the adjoints of the
-- result of the map.
vjpMap :: VjpOps -> [Adj] -> StmAux () -> SubExp -> Lambda SOACS -> [VName] -> ADM ()
vjpMap :: VjpOps
-> [Adj]
-> StmAux ()
-> SubExp
-> Lambda SOACS
-> [VName]
-> ADM ()
vjpMap VjpOps
ops [Adj]
res_adjs StmAux ()
_ SubExp
w Lambda SOACS
map_lam [VName]
as
  | Just [[(InBounds, SubExp, SubExp)]]
res_ivs <- (Adj -> Maybe [(InBounds, SubExp, SubExp)])
-> [Adj] -> Maybe [[(InBounds, SubExp, SubExp)]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Adj -> Maybe [(InBounds, SubExp, SubExp)]
isSparse [Adj]
res_adjs = ADM () -> ADM ()
forall a. ADM a -> ADM a
returnSweepCode (ADM () -> ADM ()) -> ADM () -> ADM ()
forall a b. (a -> b) -> a -> b
$ do
      -- Since at most only a constant number of adjoint are nonzero
      -- (length res_ivs), there is no need for the return sweep code to
      -- contain a Map at all.

      free <- (VName -> ADM Bool) -> [VName] -> ADM [VName]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM VName -> ADM Bool
isActive ([VName] -> ADM [VName]) -> [VName] -> ADM [VName]
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
map_lam Names -> Names -> Names
`namesSubtract` [VName] -> Names
namesFromList [VName]
as
      free_ts <- mapM lookupType free
      let adjs_for = (Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
map_lam) [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
free
          adjs_ts = (Param Type -> Type) -> [Param Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
map_lam) [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
free_ts

      let oneHot Int
res_i Adj
adj_v = (Int -> Type -> Adj) -> [Int] -> [Type] -> [Adj]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Type -> Adj
f [Int
0 :: Int ..] ([Type] -> [Adj]) -> [Type] -> [Adj]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
map_lam
            where
              f :: Int -> Type -> Adj
f Int
j Type
t
                | Int
res_i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j = Adj
adj_v
                | Bool
otherwise = Shape -> PrimType -> Adj
AdjZero (Type -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
t) (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
t)
          -- Values for the out-of-bounds case does not matter, as we will
          -- be writing to an out-of-bounds index anyway, which is ignored.
          ooBounds SubExp
adj_i = ADM (Body SOACS, [SubExp] -> [Adj])
-> ADM (Body SOACS, [SubExp] -> [Adj])
forall a. ADM a -> ADM a
subAD (ADM (Body SOACS, [SubExp] -> [Adj])
 -> ADM (Body SOACS, [SubExp] -> [Adj]))
-> (ADM (Result, [SubExp] -> [Adj])
    -> ADM (Body SOACS, [SubExp] -> [Adj]))
-> ADM (Result, [SubExp] -> [Adj])
-> ADM (Body SOACS, [SubExp] -> [Adj])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ADM (Result, [SubExp] -> [Adj])
-> ADM (Body (Rep ADM), [SubExp] -> [Adj])
ADM (Result, [SubExp] -> [Adj])
-> ADM (Body SOACS, [SubExp] -> [Adj])
forall (m :: * -> *) a.
MonadBuilder m =>
m (Result, a) -> m (Body (Rep m), a)
buildRenamedBody (ADM (Result, [SubExp] -> [Adj])
 -> ADM (Body SOACS, [SubExp] -> [Adj]))
-> ADM (Result, [SubExp] -> [Adj])
-> ADM (Body SOACS, [SubExp] -> [Adj])
forall a b. (a -> b) -> a -> b
$ do
            [(VName, Type)] -> ((VName, Type) -> ADM ()) -> ADM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName] -> [Type] -> [(VName, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
as [Type]
adjs_ts) (((VName, Type) -> ADM ()) -> ADM ())
-> ((VName, Type) -> ADM ()) -> ADM ()
forall a b. (a -> b) -> a -> b
$ \(VName
a, Type
t) -> do
              scratch <- Name -> Exp (Rep ADM) -> ADM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"oo_scratch" (Exp SOACS -> ADM SubExp) -> ADM (Exp SOACS) -> ADM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Type -> ADM (Exp (Rep ADM))
forall (m :: * -> *). MonadBuilder m => Type -> m (Exp (Rep m))
eBlank Type
t
              updateAdjIndex a (OutOfBounds, adj_i) scratch
            -- We must make sure that all free variables have the same
            -- representation in the oo-branch as in the ib-branch.
            -- In practice we do this by manifesting the adjoint.
            -- This is probably efficient, since the adjoint of a free
            -- variable is probably either a scalar or an accumulator.
            [VName] -> (VName -> ADM ()) -> ADM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [VName]
free ((VName -> ADM ()) -> ADM ()) -> (VName -> ADM ()) -> ADM ()
forall a b. (a -> b) -> a -> b
$ \VName
v -> VName -> VName -> ADM ()
insAdj VName
v (VName -> ADM ()) -> ADM VName -> ADM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Adj -> ADM VName
adjVal (Adj -> 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
            ([SubExp] -> Result)
-> ([SubExp], [SubExp] -> [Adj]) -> (Result, [SubExp] -> [Adj])
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first [SubExp] -> Result
subExpsRes (([SubExp], [SubExp] -> [Adj]) -> (Result, [SubExp] -> [Adj]))
-> ([Adj] -> ([SubExp], [SubExp] -> [Adj]))
-> [Adj]
-> (Result, [SubExp] -> [Adj])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Adj] -> ([SubExp], [SubExp] -> [Adj])
adjsReps ([Adj] -> (Result, [SubExp] -> [Adj]))
-> ADM [Adj] -> ADM (Result, [SubExp] -> [Adj])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> ADM Adj) -> [VName] -> ADM [Adj]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> ADM Adj
lookupAdj ([VName]
as [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
free)
          inBounds Int
res_i SubExp
adj_i SubExp
adj_v = ADM (Body SOACS, [SubExp] -> [Adj])
-> ADM (Body SOACS, [SubExp] -> [Adj])
forall a. ADM a -> ADM a
subAD (ADM (Body SOACS, [SubExp] -> [Adj])
 -> ADM (Body SOACS, [SubExp] -> [Adj]))
-> (ADM (Result, [SubExp] -> [Adj])
    -> ADM (Body SOACS, [SubExp] -> [Adj]))
-> ADM (Result, [SubExp] -> [Adj])
-> ADM (Body SOACS, [SubExp] -> [Adj])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ADM (Result, [SubExp] -> [Adj])
-> ADM (Body (Rep ADM), [SubExp] -> [Adj])
ADM (Result, [SubExp] -> [Adj])
-> ADM (Body SOACS, [SubExp] -> [Adj])
forall (m :: * -> *) a.
MonadBuilder m =>
m (Result, a) -> m (Body (Rep m), a)
buildRenamedBody (ADM (Result, [SubExp] -> [Adj])
 -> ADM (Body SOACS, [SubExp] -> [Adj]))
-> ADM (Result, [SubExp] -> [Adj])
-> ADM (Body SOACS, [SubExp] -> [Adj])
forall a b. (a -> b) -> a -> b
$ do
            [(Param Type, VName)] -> ((Param Type, VName) -> ADM ()) -> ADM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param Type] -> [VName] -> [(Param Type, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
map_lam) [VName]
as) (((Param Type, VName) -> ADM ()) -> ADM ())
-> ((Param Type, VName) -> ADM ()) -> ADM ()
forall a b. (a -> b) -> a -> b
$ \(Param Type
p, VName
a) -> do
              a_t <- VName -> ADM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
a
              letBindNames [paramName p] . BasicOp . Index a $
                fullSlice a_t [DimFix adj_i]
            adj_elems <-
              (Result -> [SubExp]) -> ADM Result -> ADM [SubExp]
forall a b. (a -> b) -> ADM a -> ADM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((SubExpRes -> SubExp) -> Result -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp) (ADM Result -> ADM [SubExp])
-> (Lambda SOACS -> ADM Result) -> Lambda SOACS -> ADM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body (Rep ADM) -> ADM Result
Body SOACS -> ADM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind (Body SOACS -> ADM Result)
-> (Lambda SOACS -> Body SOACS) -> Lambda SOACS -> ADM Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody
                (Lambda SOACS -> ADM [SubExp])
-> ADM (Lambda SOACS) -> ADM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VjpOps -> [Adj] -> [VName] -> Lambda SOACS -> ADM (Lambda SOACS)
vjpLambda VjpOps
ops (Int -> Adj -> [Adj]
oneHot Int
res_i (SubExp -> Adj
AdjVal SubExp
adj_v)) [VName]
adjs_for Lambda SOACS
map_lam
            let (as_adj_elems, free_adj_elems) = splitAt (length as) adj_elems
            forM_ (zip as as_adj_elems) $ \(VName
a, SubExp
a_adj_elem) ->
              VName -> (InBounds, SubExp) -> SubExp -> ADM ()
updateAdjIndex VName
a (InBounds
AssumeBounds, SubExp
adj_i) SubExp
a_adj_elem
            forM_ (zip free free_adj_elems) $ \(VName
v, SubExp
adj_se) -> do
              adj_se_v <- Name -> Exp (Rep ADM) -> ADM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"adj_v" (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
adj_se)
              insAdj v adj_se_v
            first subExpsRes . adjsReps <$> mapM lookupAdj (as <> free)

          -- Generate an iteration of the map function for every
          -- position.  This is a bit inefficient - probably we could do
          -- some deduplication.
          forPos Int
res_i (InBounds
check, SubExp
adj_i, SubExp
adj_v) = do
            adjs <-
              case InBounds
check of
                CheckBounds Maybe SubExp
b -> do
                  (obbranch, mkadjs) <- SubExp -> ADM (Body SOACS, [SubExp] -> [Adj])
ooBounds SubExp
adj_i
                  (ibbranch, _) <- inBounds res_i adj_i adj_v
                  fmap mkadjs . letTupExp' "map_adj_elem"
                    =<< eIf
                      (maybe (eDimInBounds (eSubExp w) (eSubExp adj_i)) eSubExp b)
                      (pure ibbranch)
                      (pure obbranch)
                InBounds
AssumeBounds -> do
                  (body, mkadjs) <- Int -> SubExp -> SubExp -> ADM (Body SOACS, [SubExp] -> [Adj])
inBounds Int
res_i SubExp
adj_i SubExp
adj_v
                  mkadjs . map resSubExp <$> bodyBind body
                InBounds
OutOfBounds ->
                  (VName -> ADM Adj) -> [VName] -> ADM [Adj]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> ADM Adj
lookupAdj [VName]
as

            zipWithM setAdj (as <> free) adjs

          -- Generate an iteration of the map function for every result.
          forRes Int
res_i = ((InBounds, SubExp, SubExp) -> ADM [()])
-> [(InBounds, SubExp, SubExp)] -> ADM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Int -> (InBounds, SubExp, SubExp) -> ADM [()]
forPos Int
res_i)

      zipWithM_ forRes [0 ..] res_ivs
  where
    isSparse :: Adj -> Maybe [(InBounds, SubExp, SubExp)]
isSparse (AdjSparse (Sparse Shape
shape PrimType
_ Int
vd [(InBounds, SubExp, SubExp)]
ivs)) = do
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
drop Int
vd (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape) [SubExp] -> [SubExp] -> Bool
forall a. Eq a => a -> a -> Bool
== [SubExp
w]
      [(InBounds, SubExp, SubExp)] -> Maybe [(InBounds, SubExp, SubExp)]
forall a. a -> Maybe a
Just [(InBounds, SubExp, SubExp)]
ivs
    isSparse Adj
_ =
      Maybe [(InBounds, SubExp, SubExp)]
forall a. Maybe a
Nothing
-- See Note [Adjoints of accumulators] for how we deal with
-- accumulators - it's a bit tricky here.
vjpMap VjpOps
ops [Adj]
pat_adj StmAux ()
aux SubExp
w Lambda SOACS
map_lam [VName]
as = ADM () -> ADM ()
forall a. ADM a -> ADM a
returnSweepCode (ADM () -> ADM ()) -> ADM () -> ADM ()
forall a b. (a -> b) -> a -> b
$ do
  pat_adj_vals <- [(Adj, Type)] -> ((Adj, Type) -> ADM VName) -> ADM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([Adj] -> [Type] -> [(Adj, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Adj]
pat_adj (Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
map_lam)) (((Adj, Type) -> ADM VName) -> ADM [VName])
-> ((Adj, Type) -> ADM VName) -> ADM [VName]
forall a b. (a -> b) -> a -> b
$ \(Adj
adj, Type
t) ->
    case Type
t of
      Acc {} -> Name -> Exp (Rep ADM) -> ADM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"acc_adj_rep" (Exp SOACS -> ADM VName)
-> (VName -> Exp SOACS) -> VName -> ADM VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> (VName -> BasicOp) -> VName -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape -> SubExp -> BasicOp
Replicate ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w]) (SubExp -> BasicOp) -> (VName -> SubExp) -> VName -> BasicOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var (VName -> ADM VName) -> ADM VName -> ADM VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Adj -> ADM VName
adjVal Adj
adj
      Type
_ -> VName -> ADM VName
pushAdjShape (VName -> ADM VName) -> ADM VName -> ADM VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Adj -> ADM VName
adjVal Adj
adj

  pat_adj_params <-
    mapM (newParam "map_adj_p" . rowType <=< lookupType) pat_adj_vals

  map_lam' <- renameLambda map_lam
  let param_ts = (Param Type -> Type) -> [Param Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
map_lam')
      extra = [VName] -> [Param Type] -> [(VName, Param Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
pat_adj_vals [Param Type]
pat_adj_params
      mkAdjs = ([Adj], Lambda SOACS) -> ADM ([Adj], Lambda SOACS)
forall a. a -> ADM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Param Type -> Adj) -> [Param Type] -> [Adj]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Adj
forall t. Param t -> Adj
adjFromParam [Param Type]
pat_adj_params, Lambda SOACS
map_lam')

  mapReturnSweep ops aux (w, as) map_lam' extra mkAdjs $ \[VName]
param_contribs ->
    [(Type, VName, VName)]
-> ((Type, VName, VName) -> ADM ()) -> ADM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Type] -> [VName] -> [VName] -> [(Type, VName, VName)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Type]
param_ts [VName]
as [VName]
param_contribs) (((Type, VName, VName) -> ADM ()) -> ADM ())
-> ((Type, VName, VName) -> ADM ()) -> ADM ()
forall a b. (a -> b) -> a -> b
$ \(Type
param_t, VName
a, VName
param_contrib) ->
      case Type
param_t of
        Acc {} -> SubExp -> VName -> VName -> ADM ()
freeContrib SubExp
w VName
a (VName -> ADM ()) -> ADM VName -> ADM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> ADM VName
popAdjShape VName
param_contrib -- CHECKME
        Type
_ -> VName -> VName -> ADM ()
updateAdj VName
a (VName -> ADM ()) -> ADM VName -> ADM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> ADM VName
popAdjShape VName
param_contrib

-- | Perform VJP on a FlatMap.  The 'Adj' list is the adjoints of the results of
-- the FlatMap, including the metadata results.  See Note [Adjoints of FlatMap].
vjpFlatMap ::
  VjpOps ->
  Pat Type ->
  [Adj] ->
  StmAux () ->
  SubExp ->
  ExtLambda SOACS ->
  [VName] ->
  ADM ()
vjpFlatMap :: VjpOps
-> Pat Type
-> [Adj]
-> StmAux ()
-> SubExp
-> ExtLambda SOACS
-> [VName]
-> ADM ()
vjpFlatMap VjpOps
ops Pat Type
pat [Adj]
pat_adj StmAux ()
aux SubExp
w ExtLambda SOACS
lam [VName]
as = ADM () -> ADM ()
forall a. ADM a -> ADM a
returnSweepCode (ADM () -> ADM ()) -> ADM () -> ADM ()
forall a b. (a -> b) -> a -> b
$ do
  let ((VName
_, VName
shape_arr, VName
_, VName
offset_arr), [VName]
_) = [VName] -> ((VName, VName, VName, VName), [VName])
forall a. [a] -> ((a, a, a, a), [a])
flatMapSplitMeta ([VName] -> ((VName, VName, VName, VName), [VName]))
-> [VName] -> ((VName, VName, VName, VName), [VName])
forall a b. (a -> b) -> a -> b
$ Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
pat
      ((Adj, Adj, Adj, Adj)
_, [Adj]
val_adjs) = [Adj] -> ((Adj, Adj, Adj, Adj), [Adj])
forall a. [a] -> ((a, a, a, a), [a])
flatMapSplitMeta [Adj]
pat_adj

  -- The size of this iteration's segment, and where in the concatenated
  -- results it begins.
  size_p <- Name -> Type -> ADM (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"flatmap_size_p" (Type -> ADM (Param Type)) -> Type -> ADM (Param Type)
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64
  offset_p <- newParam "flatmap_offset_p" $ Prim int64

  -- 'Left' for a nonuniform result, 'Right' for a uniform one; see Note
  -- [Adjoints of FlatMap].
  (res_adjs, uniform_adj_vals) <-
    fmap (second catMaybes) . mapAndUnzipM resAdj $
      zip (drop 1 (lambdaReturnType lam)) val_adjs

  lam' <- renameLambda lam

  let extra =
        [(VName
shape_arr, Param Type
size_p), (VName
offset_arr, Param Type
offset_p)]
          [(VName, Param Type)]
-> [(VName, Param Type)] -> [(VName, Param Type)]
forall a. [a] -> [a] -> [a]
++ [VName] -> [Param Type] -> [(VName, Param Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
uniform_adj_vals ([Either VName (Param Type)] -> [Param Type]
forall a b. [Either a b] -> [b]
rights [Either VName (Param Type)]
res_adjs)
      mkAdjs = do
        res_adjs' <- (Either VName (Param Type) -> ADM Adj)
-> [Either VName (Param Type)] -> 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 (Param Type -> Param Type -> Either VName (Param Type) -> ADM Adj
forall {dec} {dec} {t}.
Param dec -> Param dec -> Either VName (Param t) -> ADM Adj
segmentAdj Param Type
size_p Param Type
offset_p) [Either VName (Param Type)]
res_adjs
        plain_lam <- flatMapPlainLambda (Var $ paramName size_p) lam'
        pure (res_adjs', plain_lam)

  mapReturnSweep ops aux (w, as) lam' extra mkAdjs $ \[VName]
param_contribs ->
    [(VName, VName)] -> ((VName, VName) -> ADM ()) -> ADM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
as [VName]
param_contribs) (((VName, VName) -> ADM ()) -> ADM ())
-> ((VName, VName) -> ADM ()) -> ADM ()
forall a b. (a -> b) -> a -> b
$ \(VName
a, VName
param_contrib) ->
      VName -> VName -> ADM ()
updateAdj VName
a (VName -> ADM ()) -> ADM VName -> ADM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> ADM VName
popAdjShape VName
param_contrib
  where
    resAdj :: (ExtType, Adj) -> ADM (Either VName (Param Type), Maybe VName)
resAdj (ExtType
t, Adj
adj) = do
      adj_v <- Adj -> ADM VName
adjVal Adj
adj
      if flatMapNonuniform t
        then pure (Left adj_v, Nothing)
        else do
          adj_v' <- pushAdjShape adj_v
          adj_p <- newParam "flatmap_res_adj_p" . rowType =<< lookupType adj_v'
          pure (Right adj_p, Just adj_v')

    segmentAdj :: Param dec -> Param dec -> Either VName (Param t) -> ADM Adj
segmentAdj Param dec
_ Param dec
_ (Right Param t
adj_p) = 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
$ Param t -> Adj
forall t. Param t -> Adj
adjFromParam Param t
adj_p
    segmentAdj Param dec
size_p Param dec
offset_p (Left VName
adj_v) = do
      adj_v_t <- VName -> ADM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
adj_v
      let segment =
            SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
DimSlice (VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param dec -> VName
forall dec. Param dec -> VName
paramName Param dec
offset_p) (VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param dec -> VName
forall dec. Param dec -> VName
paramName Param dec
size_p) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
      slice <- vecSlice adj_v_t [segment]
      fmap adjFromVar . letExp (baseName adj_v <> "_slice") . BasicOp $
        Index adj_v slice

-- Note [Adjoints of FlatMap]
--
-- The return sweep of a FlatMap is an ordinary Map, not a FlatMap.  Iteration
-- 'j' of the FlatMap contributed the segment of each nonuniform (concatenated)
-- result that begins at 'offset[j]' and has length 'shape[j]', where the offset
-- and shape arrays are results of the forward sweep.  This is all we need to
-- find the part of an adjoint that a given iteration is responsible for.
--
-- The inputs are necessarily regular arrays of the same length as the Map, and
-- so are handled exactly as for a Map.  Similarly, free variables of the lambda
-- receive contributions from every iteration, and so are handled with the usual
-- accumulator machinery; see Note [Adjoints of accumulators].  The metadata
-- results are integers, so their adjoints are ignored.
--
-- The adjoints of the results are provided to the lambda in one of two ways:
--
--  * A uniform result has exactly one element per iteration, so its adjoint is
--    an ordinary Map input.
--
--  * The adjoint of a nonuniform result is the concatenation of the adjoints of
--    every segment, which cannot be split into one Map input per iteration
--    because the segments have different sizes.  Instead the entire array is
--    passed as a free variable and sliced inside the lambda.
--
-- Accumulators cannot occur among the inputs or results: the lambda would have
-- to return the updated accumulator, and an accumulator can be neither
-- concatenated nor collected into an array. (This is distinct from an array
-- whose *adjoint* is an accumulator, which happens all the time.) FIXME: it is
-- possible that one way we may return adjoints as the uniform result, but this
-- should be easy to add.
--
-- # Coercing the segment size
--
-- One wrinkle remains. Inside the lambda, a nonuniform result has type '[k]t',
-- where 'k' is the size the lambda itself computes and returns as its first
-- result. The adjoint we slice out of the concatenated adjoint has size
-- 'shape[j]' instead. These are dynamically equal, but they are distinct names
-- as far as the type checker is concerned, and the adjoint of a value of type
-- '[k]t' must have type '[k]t'. We handle this by converting the ExtLambda into
-- a Lambda whose nonuniform results are coerced to size 'shape[j]' (see
-- 'flatMapPlainLambda'), and differentiate that. Everything downstream is then
-- the ordinary machinery for differentiating a lambda, and the size mismatch is
-- dealt with by the existing rule for differentiating a coercion.