    {-# LANGUAGE TypeFamilies #-}

    module Futhark.AD.Rev.Scatter (vjpScatter) where

    import Control.Monad
    import Futhark.AD.Rev.Monad
    import Futhark.Analysis.PrimExp.Convert
    import Futhark.Builder
    import Futhark.IR.SOACS
    import Futhark.Tools
    import Futhark.Util (chunk)

    withinBounds :: [(SubExp, VName)] -> TPrimExp Bool VName
    withinBounds :: [(SubExp, VName)] -> TPrimExp Bool VName
    withinBounds [] = PrimExp VName -> TPrimExp Bool VName
    forall {k} (t :: k) v. PrimExp v -> TPrimExp t v
    TPrimExp (PrimExp VName -> TPrimExp Bool VName)
    -> PrimExp VName -> TPrimExp Bool VName
    forall a b. (a -> b) -> a -> b
    $ PrimValue -> PrimExp VName
    forall v. PrimValue -> PrimExp v
    ValueExp (Bool -> PrimValue
    BoolValue Bool
    True)
    withinBounds [(SubExp
    q, VName
    i)] = (VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    i TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
    forall {k} v (t :: k).
    Eq v =>
    TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
    .<. SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    q) TPrimExp Bool VName -> TPrimExp Bool VName -> TPrimExp Bool VName
    forall v.
    Eq v =>
    TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
    .&&. (SubExp -> TPrimExp Int64 VName
    pe64 (IntType -> Integer -> SubExp
    intConst IntType
    Int64 (-Integer
    1)) TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
    forall {k} v (t :: k).
    Eq v =>
    TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
    .<. VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    i)
    withinBounds ((SubExp, VName)
    qi : [(SubExp, VName)]
    qis) = [(SubExp, VName)] -> TPrimExp Bool VName
    withinBounds [(SubExp, VName)
    qi] TPrimExp Bool VName -> TPrimExp Bool VName -> TPrimExp Bool VName
    forall v.
    Eq v =>
    TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
    .&&. [(SubExp, VName)] -> TPrimExp Bool VName
    withinBounds [(SubExp, VName)]
    qis

    -- Generates a potential tower-of-maps lambda body for an indexing operation.
    -- Assuming parameters:
    --   `arr`   the array that is indexed
    --   `[(w_1, i_1), (w_2, i_2), ..., (w_k, i_k)]` outer lambda formal parameters and their bounds
    --   `[n_1,n_2,...]ptp` the type of the index expression `arr[i_1,i_2,...,i_k]`
    -- Generates something like:
    -- (\ i_1 i_2 ->
    --    map (\j_1 -> ... if (i_1 >= 0 && i_1 < w_1) &&
    --                        (i_2 >= 0 && i_2 < w_2) && ...
    --                     then arr[i_1, i_2, ... j_1, ...]
    --                     else 0
    --        ) (iota n_1)
    -- )
    -- The idea is that you do not want to put under the `if` something
    --     that is an array because it would not flatten well!
    genIdxLamBody :: VName -> [(SubExp, Param Type)] -> Type -> ADM (Body SOACS)
    genIdxLamBody :: VName -> [(SubExp, Param Type)] -> Type -> ADM (Body SOACS)
    genIdxLamBody VName
    as [(SubExp, Param Type)]
    wpis = VName
    -> [(SubExp, Param Type)]
    -> [Param Type]
    -> Type
    -> ADM (Body SOACS)
    genRecLamBody VName
    as [(SubExp, Param Type)]
    wpis []
      where
        genRecLamBody :: VName -> [(SubExp, Param Type)] -> [Param Type] -> Type -> ADM (Body SOACS)
        genRecLamBody :: VName
    -> [(SubExp, Param Type)]
    -> [Param Type]
    -> Type
    -> ADM (Body SOACS)
    genRecLamBody VName
    arr [(SubExp, Param Type)]
    w_pis [Param Type]
    nest_pis (Array PrimType
    t (Shape []) NoUniqueness
    _) =
          VName
    -> [(SubExp, Param Type)]
    -> [Param Type]
    -> Type
    -> ADM (Body SOACS)
    genRecLamBody VName
    arr [(SubExp, Param Type)]
    w_pis [Param Type]
    nest_pis (PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    t)
        genRecLamBody VName
    arr [(SubExp, Param Type)]
    w_pis [Param Type]
    nest_pis (Array PrimType
    t (Shape (SubExp
    s : [SubExp]
    ss)) NoUniqueness
    _) = do
          Param Type
    new_ip <- [Char] -> Type -> ADM (Param Type)
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    [Char] -> dec -> m (Param dec)
    newParam [Char]
    "i" (PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64)
          let t' :: Type
    t' = PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    t Type -> ShapeBase SubExp -> Type
    `arrayOfShape` [SubExp] -> ShapeBase SubExp
    forall d. [d] -> ShapeBase d
    Shape [SubExp]
    ss
          Lambda SOACS
    inner_lam <-
            [LParam (Rep ADM)] -> ADM Result -> ADM (Lambda (Rep ADM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
    mkLambda [Param Type
    LParam (Rep ADM)
    new_ip] (ADM Result -> ADM (Lambda (Rep ADM)))
    -> ADM Result -> ADM (Lambda (Rep ADM))
    forall a b. (a -> b) -> a -> b
    $
              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
    =<< VName
    -> [(SubExp, Param Type)]
    -> [Param Type]
    -> Type
    -> ADM (Body SOACS)
    genRecLamBody VName
    arr [(SubExp, Param Type)]
    w_pis ([Param Type]
    nest_pis [Param Type] -> [Param Type] -> [Param Type]
    forall a. [a] -> [a] -> [a]
    ++ [Param Type
    new_ip]) Type
    t'
          let ([SubExp]
    _, [Param Type]
    orig_pis) = [(SubExp, Param Type)] -> ([SubExp], [Param Type])
    forall a b. [(a, b)] -> ([a], [b])
    unzip [(SubExp, Param Type)]
    w_pis
          ADM Result -> ADM (Body (Rep ADM))
    ADM Result -> ADM (Body SOACS)
    forall (m :: * -> *).
    MonadBuilder m =>
    m Result -> m (Body (Rep m))
    buildBody_ (ADM Result -> ADM (Body SOACS))
    -> (ADM Result -> ADM Result) -> ADM Result -> ADM (Body SOACS)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Scope SOACS -> ADM Result -> ADM Result
    forall a. Scope SOACS -> ADM a -> ADM a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope ([Param Type] -> Scope SOACS
    forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfLParams ([Param Type]
    orig_pis [Param Type] -> [Param Type] -> [Param Type]
    forall a. [a] -> [a] -> [a]
    ++ [Param Type]
    nest_pis)) (ADM Result -> ADM (Body SOACS)) -> ADM Result -> ADM (Body SOACS)
    forall a b. (a -> b) -> a -> b
    $ do
            VName
    iota_v <- [Char] -> Exp (Rep ADM) -> ADM VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "iota" (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 -> SubExp -> SubExp -> IntType -> BasicOp
    Iota SubExp
    s (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    0) (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1) IntType
    Int64
            SubExp
    r <- [Char] -> Exp (Rep ADM) -> ADM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp (VName -> [Char]
    baseString VName
    arr [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_elem") (Exp (Rep ADM) -> ADM SubExp) -> Exp (Rep ADM) -> ADM SubExp
    forall a b. (a -> b) -> a -> b
    $ Op (Rep ADM) -> Exp (Rep ADM)
    forall rep. Op rep -> Exp rep
    Op (Op (Rep ADM) -> Exp (Rep ADM)) -> Op (Rep ADM) -> Exp (Rep ADM)
    forall a b. (a -> b) -> a -> b
    $ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
    forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
    Screma SubExp
    s [VName
    iota_v] (Lambda SOACS -> ScremaForm SOACS
    forall rep. Lambda rep -> ScremaForm rep
    mapSOAC Lambda SOACS
    inner_lam)
            Result -> ADM Result
    forall a. a -> ADM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [SubExp -> SubExpRes
    subExpRes SubExp
    r]
        genRecLamBody VName
    arr [(SubExp, Param Type)]
    w_pis [Param Type]
    nest_pis (Prim PrimType
    ptp) = do
          let ([SubExp]
    ws, [Param Type]
    orig_pis) = [(SubExp, Param Type)] -> ([SubExp], [Param Type])
    forall a b. [(a, b)] -> ([a], [b])
    unzip [(SubExp, Param Type)]
    w_pis
          let inds :: [VName]
    inds = (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]
    orig_pis [Param Type] -> [Param Type] -> [Param Type]
    forall a. [a] -> [a] -> [a]
    ++ [Param Type]
    nest_pis)
          Scope SOACS -> ADM (Body SOACS) -> ADM (Body SOACS)
    forall a. Scope SOACS -> ADM a -> ADM a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope ([Param Type] -> Scope SOACS
    forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfLParams ([Param Type]
    orig_pis [Param Type] -> [Param Type] -> [Param Type]
    forall a. [a] -> [a] -> [a]
    ++ [Param Type]
    nest_pis)) (ADM (Body SOACS) -> ADM (Body SOACS))
    -> ADM (Body SOACS) -> ADM (Body SOACS)
    forall a b. (a -> b) -> a -> b
    $
            [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))
    -> 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
                  (TPrimExp Bool VName -> ADM (Exp (Rep ADM))
    forall a (m :: * -> *).
    (ToExp a, MonadBuilder m) =>
    a -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    TPrimExp Bool VName -> m (Exp (Rep m))
    toExp (TPrimExp Bool VName -> ADM (Exp (Rep ADM)))
    -> TPrimExp Bool VName -> ADM (Exp (Rep ADM))
    forall a b. (a -> b) -> a -> b
    $ [(SubExp, VName)] -> TPrimExp Bool VName
    withinBounds ([(SubExp, VName)] -> TPrimExp Bool VName)
    -> [(SubExp, VName)] -> TPrimExp Bool VName
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> [VName] -> [(SubExp, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [SubExp]
    ws ([VName] -> [(SubExp, VName)]) -> [VName] -> [(SubExp, VName)]
    forall a b. (a -> b) -> a -> b
    $ (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]
    orig_pis)
                  ( do
                      SubExp
    r <- [Char] -> Exp (Rep ADM) -> ADM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "r" (Exp (Rep ADM) -> ADM SubExp) -> Exp (Rep ADM) -> ADM SubExp
    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
    arr (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ [DimIndex SubExp] -> Slice SubExp
    forall d. [DimIndex d] -> Slice d
    Slice ([DimIndex SubExp] -> Slice SubExp)
    -> [DimIndex SubExp] -> Slice SubExp
    forall a b. (a -> b) -> a -> b
    $ (VName -> DimIndex SubExp) -> [VName] -> [DimIndex SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map (SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix (SubExp -> DimIndex SubExp)
    -> (VName -> SubExp) -> VName -> DimIndex SubExp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> SubExp
    Var) [VName]
    inds
                      [SubExp] -> ADM (Body (Rep ADM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [SubExp
    r]
                  )
                  ([SubExp] -> ADM (Body (Rep ADM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [PrimValue -> SubExp
    Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
    forall a b. (a -> b) -> a -> b
    $ PrimType -> PrimValue
    blankPrimValue PrimType
    ptp])
              ]
        genRecLamBody VName
    _ [(SubExp, Param Type)]
    _ [Param Type]
    _ Type
    _ = [Char] -> ADM (Body SOACS)
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "In Rev.hs, helper function genRecLamBody, unreachable case reached!"

    --
    -- Original:
    --   let ys = scatter xs is vs
    -- Assumes no duplicate indices in `is`
    -- Forward Sweep:
    --   let xs_save = gather xs is
    --   let ys = scatter xs is vs
    -- Return Sweep:
    --   let vs_ctrbs = gather is ys_adj
    --   let vs_adj \overline{+}= vs_ctrbs -- by map or generalized reduction
    --   let xs_adj = scatter ys_adj is \overline{0}
    --   let xs = scatter ys is xs_save
    vjpScatter1 ::
      PatElem Type ->
      StmAux () ->
      (SubExp, [VName], (ShapeBase SubExp, Int, VName)) ->
      ADM () ->
      ADM ()
    vjpScatter1 :: PatElem Type
    -> StmAux ()
    -> (SubExp, [VName], (ShapeBase SubExp, Int, VName))
    -> ADM ()
    -> ADM ()
    vjpScatter1 PatElem Type
    pys StmAux ()
    aux (SubExp
    w, [VName]
    ass, (ShapeBase SubExp
    shp, Int
    num_vals, VName
    xs)) ADM ()
    m = do
      let rank :: Int
    rank = [SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length ([SubExp] -> Int) -> [SubExp] -> Int
    forall a b. (a -> b) -> a -> b
    $ ShapeBase SubExp -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims ShapeBase SubExp
    shp
          ([VName]
    all_inds, [VName]
    val_as) = Int -> [VName] -> ([VName], [VName])
    forall a. Int -> [a] -> ([a], [a])
    splitAt (Int
    rank Int -> Int -> Int
    forall a. Num a => a -> a -> a
    * Int
    num_vals) [VName]
    ass
          inds_as :: [[VName]]
    inds_as = Int -> [VName] -> [[VName]]
    forall a. Int -> [a] -> [[a]]
    chunk Int
    rank [VName]
    all_inds
      Type
    xs_t <- VName -> ADM Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType VName
    xs
      let val_t :: Type
    val_t = Int -> Type -> Type
    forall u.
    Int
    -> TypeBase (ShapeBase SubExp) u -> TypeBase (ShapeBase SubExp) u
    stripArray (ShapeBase SubExp -> Int
    forall a. ArrayShape a => a -> Int
    shapeRank ShapeBase SubExp
    shp) Type
    xs_t
      -- computing xs_save
      [VName]
    xs_saves <- [[VName]] -> VName -> Type -> ADM [VName]
    mkGather [[VName]]
    inds_as VName
    xs Type
    xs_t
      -- performing the scatter
      Lambda SOACS
    id_lam <-
        [Type] -> ADM (Lambda SOACS)
    forall rep (m :: * -> *).
    (Buildable rep, MonadFreshNames m) =>
    [Type] -> m (Lambda rep)
    mkIdentityLambda ([Type] -> ADM (Lambda SOACS)) -> [Type] -> ADM (Lambda SOACS)
    forall a b. (a -> b) -> a -> b
    $
          Int -> Type -> [Type]
    forall a. Int -> a -> [a]
    replicate (ShapeBase SubExp -> Int
    forall a. ArrayShape a => a -> Int
    shapeRank ShapeBase SubExp
    shp) (PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64) [Type] -> [Type] -> [Type]
    forall a. [a] -> [a] -> [a]
    ++ Int -> Type -> [Type]
    forall a. Int -> a -> [a]
    replicate (ShapeBase SubExp -> Int
    forall a. ArrayShape a => a -> Int
    shapeRank ShapeBase SubExp
    shp) Type
    val_t
      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 ([PatElem Type] -> Pat Type
    forall dec. [PatElem dec] -> Pat dec
    Pat [PatElem Type
    pys]) 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
    $ Op (Rep ADM) -> Exp (Rep ADM)
    forall rep. Op rep -> Exp rep
    Op (Op (Rep ADM) -> Exp (Rep ADM)) -> Op (Rep ADM) -> Exp (Rep ADM)
    forall a b. (a -> b) -> a -> b
    $ SubExp
    -> [VName] -> ScatterSpec VName -> Lambda SOACS -> SOAC SOACS
    forall rep.
    SubExp -> [VName] -> ScatterSpec VName -> Lambda rep -> SOAC rep
    Scatter SubExp
    w [VName]
    ass [(ShapeBase SubExp
    shp, Int
    num_vals, VName
    xs)] Lambda SOACS
    id_lam
      ADM ()
    m
      let ys :: VName
    ys = PatElem Type -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem Type
    pys
      -- XXX: Since our restoration of xs will consume ys, we have to
      -- make a copy of ys in the chance that it is actually the result
      -- of the program.  In that case the asymptotics will not be
      -- (locally) preserved, but since ys must necessarily have been
      -- constructed somewhere close, they are probably globally OK.
      VName
    ys_copy <-
        [Char] -> Exp (Rep ADM) -> ADM VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp (VName -> [Char]
    baseString VName
    ys [Char] -> [Char] -> [Char]
    forall a. Semigroup a => a -> a -> a
    <> [Char]
    "_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
    $
          ShapeBase SubExp -> SubExp -> BasicOp
    Replicate ShapeBase SubExp
    forall a. Monoid a => a
    mempty (VName -> SubExp
    Var VName
    ys)
      ADM () -> ADM ()
    forall a. ADM a -> ADM a
    returnSweepCode (ADM () -> ADM ()) -> ADM () -> ADM ()
    forall a b. (a -> b) -> a -> b
    $ do
        VName
    ys_adj <- VName -> ADM VName
    lookupAdjVal VName
    ys
        -- computing vs_ctrbs and updating vs_adj
        [VName]
    vs_ctrbs <- [[VName]] -> VName -> Type -> ADM [VName]
    mkGather [[VName]]
    inds_as VName
    ys_adj Type
    xs_t
        (VName -> VName -> ADM ()) -> [VName] -> [VName] -> ADM ()
    forall (m :: * -> *) a b c.
    Applicative m =>
    (a -> b -> m c) -> [a] -> [b] -> m ()
    zipWithM_ VName -> VName -> ADM ()
    updateAdj [VName]
    val_as [VName]
    vs_ctrbs -- use Slice?
        -- creating xs_adj
        [VName]
    zeros <-
          Int -> ADM VName -> ADM [VName]
    forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
    replicateM ([VName] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [VName]
    val_as) (ADM VName -> ADM [VName])
    -> (Exp SOACS -> ADM VName) -> Exp SOACS -> ADM [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [Char] -> Exp (Rep ADM) -> ADM VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "zeros" (Exp SOACS -> ADM [VName]) -> Exp SOACS -> ADM [VName]
    forall a b. (a -> b) -> a -> b
    $
            Type -> Exp SOACS
    forall rep. Type -> Exp rep
    zeroExp (Type -> Exp SOACS) -> Type -> Exp SOACS
    forall a b. (a -> b) -> a -> b
    $
              Type
    xs_t Type -> SubExp -> Type
    forall d u.
    ArrayShape (ShapeBase d) =>
    TypeBase (ShapeBase d) u -> d -> TypeBase (ShapeBase d) u
    `setOuterSize` SubExp
    w
        let f_tps :: [Type]
    f_tps = Int -> Type -> [Type]
    forall a. Int -> a -> [a]
    replicate (Int
    rank Int -> Int -> Int
    forall a. Num a => a -> a -> a
    * Int
    num_vals) (PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64) [Type] -> [Type] -> [Type]
    forall a. [a] -> [a] -> [a]
    ++ Int -> Type -> [Type]
    forall a. Int -> a -> [a]
    replicate Int
    num_vals Type
    val_t
        Lambda SOACS
    f <- [Type] -> ADM (Lambda SOACS)
    forall rep (m :: * -> *).
    (Buildable rep, MonadFreshNames m) =>
    [Type] -> m (Lambda rep)
    mkIdentityLambda [Type]
    f_tps
        VName
    xs_adj <-
          [Char] -> Exp (Rep ADM) -> ADM VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp (VName -> [Char]
    baseString VName
    xs [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_adj") (Exp SOACS -> ADM VName)
    -> (SOAC SOACS -> Exp SOACS) -> SOAC SOACS -> ADM VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Op SOACS -> Exp SOACS
    SOAC SOACS -> Exp SOACS
    forall rep. Op rep -> Exp rep
    Op (SOAC SOACS -> ADM VName) -> SOAC SOACS -> ADM VName
    forall a b. (a -> b) -> a -> b
    $
            SubExp
    -> [VName] -> ScatterSpec VName -> Lambda SOACS -> SOAC SOACS
    forall rep.
    SubExp -> [VName] -> ScatterSpec VName -> Lambda rep -> SOAC rep
    Scatter SubExp
    w ([VName]
    all_inds [VName] -> [VName] -> [VName]
    forall a. [a] -> [a] -> [a]
    ++ [VName]
    zeros) [(ShapeBase SubExp
    shp, Int
    num_vals, VName
    ys_adj)] Lambda SOACS
    f
        VName -> VName -> ADM ()
    insAdj VName
    xs VName
    xs_adj -- reusing the ys_adj for xs_adj!
        Lambda SOACS
    f' <- [Type] -> ADM (Lambda SOACS)
    forall rep (m :: * -> *).
    (Buildable rep, MonadFreshNames m) =>
    [Type] -> m (Lambda rep)
    mkIdentityLambda [Type]
    f_tps
        VName
    xs_rc <-
          StmAux () -> ADM VName -> ADM VName
    forall (m :: * -> *) anyrep a.
    MonadBuilder m =>
    StmAux anyrep -> m a -> m a
    auxing StmAux ()
    aux (ADM VName -> ADM VName)
    -> (SOAC SOACS -> ADM VName) -> SOAC SOACS -> ADM VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [Char] -> Exp (Rep ADM) -> ADM VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp (VName -> [Char]
    baseString VName
    xs [Char] -> [Char] -> [Char]
    forall a. Semigroup a => a -> a -> a
    <> [Char]
    "_rc") (Exp SOACS -> ADM VName)
    -> (SOAC SOACS -> Exp SOACS) -> SOAC SOACS -> ADM VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Op SOACS -> Exp SOACS
    SOAC SOACS -> Exp SOACS
    forall rep. Op rep -> Exp rep
    Op (SOAC SOACS -> ADM VName) -> SOAC SOACS -> ADM VName
    forall a b. (a -> b) -> a -> b
    $
            SubExp
    -> [VName] -> ScatterSpec VName -> Lambda SOACS -> SOAC SOACS
    forall rep.
    SubExp -> [VName] -> ScatterSpec VName -> Lambda rep -> SOAC rep
    Scatter SubExp
    w ([VName]
    all_inds [VName] -> [VName] -> [VName]
    forall a. [a] -> [a] -> [a]
    ++ [VName]
    xs_saves) [(ShapeBase SubExp
    shp, Int
    num_vals, VName
    ys)] Lambda SOACS
    f'
        VName -> VName -> ADM ()
    addSubstitution VName
    xs VName
    xs_rc
        VName -> VName -> ADM ()
    addSubstitution VName
    ys VName
    ys_copy
      where
        -- Creates a potential map-nest that indexes in full the array,
        --   and applies the condition of indices within bounds at the
        --   deepest level in the nest so that everything can be parallel.
        mkGather :: [[VName]] -> VName -> Type -> ADM [VName]
        mkGather :: [[VName]] -> VName -> Type -> ADM [VName]
    mkGather [[VName]]
    inds_as VName
    arr Type
    arr_t = do
          [[Param Type]]
    ips <- [[VName]] -> ([VName] -> ADM [Param Type]) -> ADM [[Param Type]]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [[VName]]
    inds_as (([VName] -> ADM [Param Type]) -> ADM [[Param Type]])
    -> ([VName] -> ADM [Param Type]) -> ADM [[Param Type]]
    forall a b. (a -> b) -> a -> b
    $ \[VName]
    idxs ->
            (VName -> ADM (Param Type)) -> [VName] -> ADM [Param Type]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (\VName
    idx -> [Char] -> Type -> ADM (Param Type)
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    [Char] -> dec -> m (Param dec)
    newParam (VName -> [Char]
    baseString VName
    idx [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_elem") (PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64)) [VName]
    idxs

          Lambda SOACS
    gather_lam <- [LParam (Rep ADM)] -> ADM Result -> ADM (Lambda (Rep ADM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
    mkLambda ([[Param Type]] -> [Param Type]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [[Param Type]]
    ips) (ADM Result -> ADM (Lambda SOACS))
    -> (([Param Type] -> ADM Result) -> ADM Result)
    -> ([Param Type] -> ADM Result)
    -> ADM (Lambda SOACS)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ([Result] -> Result) -> ADM [Result] -> ADM Result
    forall a b. (a -> b) -> ADM a -> ADM b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [Result] -> Result
    forall a. Monoid a => [a] -> a
    mconcat (ADM [Result] -> ADM Result)
    -> (([Param Type] -> ADM Result) -> ADM [Result])
    -> ([Param Type] -> ADM Result)
    -> ADM Result
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [[Param Type]] -> ([Param Type] -> ADM Result) -> ADM [Result]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [[Param Type]]
    ips (([Param Type] -> ADM Result) -> ADM (Lambda SOACS))
    -> ([Param Type] -> ADM Result) -> ADM (Lambda SOACS)
    forall a b. (a -> b) -> a -> b
    $ \[Param Type]
    idxs -> do
            let q :: Int
    q = [Param Type] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Param Type]
    idxs
                ([SubExp]
    ws, Type
    eltp) = (Int -> [SubExp] -> [SubExp]
    forall a. Int -> [a] -> [a]
    take Int
    q ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ Type -> [SubExp]
    forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
    arrayDims Type
    arr_t, Int -> Type -> Type
    forall u.
    Int
    -> TypeBase (ShapeBase SubExp) u -> TypeBase (ShapeBase SubExp) u
    stripArray Int
    q Type
    arr_t)
            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
    =<< VName -> [(SubExp, Param Type)] -> Type -> ADM (Body SOACS)
    genIdxLamBody VName
    arr ([SubExp] -> [Param Type] -> [(SubExp, Param Type)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [SubExp]
    ws [Param Type]
    idxs) Type
    eltp
          let soac :: SOAC SOACS
    soac = SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
    forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
    Screma SubExp
    w ([[VName]] -> [VName]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [[VName]]
    inds_as) (Lambda SOACS -> ScremaForm SOACS
    forall rep. Lambda rep -> ScremaForm rep
    mapSOAC Lambda SOACS
    gather_lam)
          [Char] -> Exp (Rep ADM) -> ADM [VName]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [VName]
    letTupExp (VName -> [Char]
    baseString VName
    arr [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_gather") (Exp (Rep ADM) -> ADM [VName]) -> Exp (Rep ADM) -> ADM [VName]
    forall a b. (a -> b) -> a -> b
    $ Op SOACS -> Exp SOACS
    forall rep. Op rep -> Exp rep
    Op Op SOACS
    SOAC SOACS
    soac

    vjpScatter ::
      VjpOps ->
      Pat Type ->
      StmAux () ->
      (SubExp, [VName], Lambda SOACS, [(Shape, Int, VName)]) ->
      ADM () ->
      ADM ()
    vjpScatter :: VjpOps
    -> Pat Type
    -> StmAux ()
    -> (SubExp, [VName], Lambda SOACS, ScatterSpec VName)
    -> ADM ()
    -> ADM ()
    vjpScatter VjpOps
    ops (Pat [PatElem Type]
    pes) StmAux ()
    aux (SubExp
    w, [VName]
    ass, Lambda SOACS
    lam, ScatterSpec VName
    written_info) ADM ()
    m
      | Lambda SOACS -> Bool
    forall rep. Lambda rep -> Bool
    isIdentityLambda Lambda SOACS
    lam,
        [(ShapeBase SubExp
    shp, Int
    num_vals, VName
    xs)] <- ScatterSpec VName
    written_info,
        [PatElem Type
    pys] <- [PatElem Type]
    pes =
          PatElem Type
    -> StmAux ()
    -> (SubExp, [VName], (ShapeBase SubExp, Int, VName))
    -> ADM ()
    -> ADM ()
    vjpScatter1 PatElem Type
    pys StmAux ()
    aux (SubExp
    w, [VName]
    ass, (ShapeBase SubExp
    shp, Int
    num_vals, VName
    xs)) ADM ()
    m
      | Lambda SOACS -> Bool
    forall rep. Lambda rep -> Bool
    isIdentityLambda Lambda SOACS
    lam = do
          let sind :: Int
    sind = ScatterSpec VName -> Int
    forall {a} {c}. [(ShapeBase a, Int, c)] -> Int
    splitInd ScatterSpec VName
    written_info
              ([VName]
    inds, [VName]
    vals) = Int -> [VName] -> ([VName], [VName])
    forall a. Int -> [a] -> ([a], [a])
    splitAt Int
    sind [VName]
    ass
          [Stm SOACS]
    lst_stms <- ([VName], [VName])
    -> [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    -> ADM [Stm SOACS]
    chunkScatterInps ([VName]
    inds, [VName]
    vals) ([PatElem Type]
    -> ScatterSpec VName
    -> [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [PatElem Type]
    pes ScatterSpec VName
    written_info)
          Stms SOACS -> ADM ()
    diffScatters ([Stm SOACS] -> Stms SOACS
    forall rep. [Stm rep] -> Stms rep
    stmsFromList [Stm SOACS]
    lst_stms)
      | Bool
    otherwise =
          [Char] -> ADM ()
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "vjpScatter: cannot handle"
      where
        splitInd :: [(ShapeBase a, Int, c)] -> Int
    splitInd [] = Int
    0
        splitInd ((ShapeBase a
    shp, Int
    num_res, c
    _) : [(ShapeBase a, Int, c)]
    rest) =
          Int
    num_res Int -> Int -> Int
    forall a. Num a => a -> a -> a
    * [a] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length (ShapeBase a -> [a]
    forall d. ShapeBase d -> [d]
    shapeDims ShapeBase a
    shp) Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + [(ShapeBase a, Int, c)] -> Int
    splitInd [(ShapeBase a, Int, c)]
    rest
        chunkScatterInps :: ([VName], [VName])
    -> [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    -> ADM [Stm SOACS]
    chunkScatterInps ([VName]
    acc_inds, [VName]
    acc_vals) [] =
          case ([VName]
    acc_inds, [VName]
    acc_vals) of
            ([], []) -> [Stm SOACS] -> ADM [Stm SOACS]
    forall a. a -> ADM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure []
            ([VName], [VName])
    _ -> [Char] -> ADM [Stm SOACS]
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "chunkScatterInps: cannot handle"
        chunkScatterInps
          ([VName]
    acc_inds, [VName]
    acc_vals)
          ((PatElem Type
    pe, info :: (ShapeBase SubExp, Int, VName)
    info@(ShapeBase SubExp
    shp, Int
    num_vals, VName
    _)) : [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    rest) = do
            let num_inds :: Int
    num_inds = Int
    num_vals Int -> Int -> Int
    forall a. Num a => a -> a -> a
    * [SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length (ShapeBase SubExp -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims ShapeBase SubExp
    shp)
                ([VName]
    curr_inds, [VName]
    other_inds) = Int -> [VName] -> ([VName], [VName])
    forall a. Int -> [a] -> ([a], [a])
    splitAt Int
    num_inds [VName]
    acc_inds
                ([VName]
    curr_vals, [VName]
    other_vals) = Int -> [VName] -> ([VName], [VName])
    forall a. Int -> [a] -> ([a], [a])
    splitAt Int
    num_vals [VName]
    acc_vals
            [Type]
    vtps <- (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 VName -> ADM Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType [VName]
    curr_vals
            Lambda SOACS
    f <- [Type] -> ADM (Lambda SOACS)
    forall rep (m :: * -> *).
    (Buildable rep, MonadFreshNames m) =>
    [Type] -> m (Lambda rep)
    mkIdentityLambda (Int -> Type -> [Type]
    forall a. Int -> a -> [a]
    replicate Int
    num_inds (PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64) [Type] -> [Type] -> [Type]
    forall a. [a] -> [a] -> [a]
    ++ [Type]
    vtps)
            let stm :: Stm SOACS
    stm =
                  Pat (LetDec SOACS)
    -> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let ([PatElem Type] -> Pat Type
    forall dec. [PatElem dec] -> Pat dec
    Pat [PatElem Type
    pe]) StmAux ()
    StmAux (ExpDec SOACS)
    aux (Exp SOACS -> Stm SOACS)
    -> (SOAC SOACS -> Exp SOACS) -> SOAC SOACS -> Stm SOACS
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Op SOACS -> Exp SOACS
    SOAC SOACS -> Exp SOACS
    forall rep. Op rep -> Exp rep
    Op (SOAC SOACS -> Stm SOACS) -> SOAC SOACS -> Stm SOACS
    forall a b. (a -> b) -> a -> b
    $
                    SubExp
    -> [VName] -> ScatterSpec VName -> Lambda SOACS -> SOAC SOACS
    forall rep.
    SubExp -> [VName] -> ScatterSpec VName -> Lambda rep -> SOAC rep
    Scatter SubExp
    w ([VName]
    curr_inds [VName] -> [VName] -> [VName]
    forall a. [a] -> [a] -> [a]
    ++ [VName]
    curr_vals) [(ShapeBase SubExp, Int, VName)
    info] Lambda SOACS
    f
            [Stm SOACS]
    stms_rest <- ([VName], [VName])
    -> [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    -> ADM [Stm SOACS]
    chunkScatterInps ([VName]
    other_inds, [VName]
    other_vals) [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    rest
            [Stm SOACS] -> ADM [Stm SOACS]
    forall a. a -> ADM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([Stm SOACS] -> ADM [Stm SOACS]) -> [Stm SOACS] -> ADM [Stm SOACS]
    forall a b. (a -> b) -> a -> b
    $ Stm SOACS
    stm Stm SOACS -> [Stm SOACS] -> [Stm SOACS]
    forall a. a -> [a] -> [a]
    : [Stm SOACS]
    stms_rest
        diffScatters :: Stms SOACS -> ADM ()
    diffScatters Stms SOACS
    all_stms
          | Just (Stm SOACS
    stm, Stms SOACS
    stms) <- Stms SOACS -> Maybe (Stm SOACS, Stms SOACS)
    forall rep. Stms rep -> Maybe (Stm rep, Stms rep)
    stmsHead Stms SOACS
    all_stms =
              VjpOps -> Stm SOACS -> ADM () -> ADM ()
    vjpStm VjpOps
    ops Stm SOACS
    stm (ADM () -> ADM ()) -> ADM () -> ADM ()
    forall a b. (a -> b) -> a -> b
    $ Stms SOACS -> ADM ()
    diffScatters Stms SOACS
    stms
          | Bool
    otherwise = ADM ()
    m