    {-# LANGUAGE TypeFamilies #-}

    -- | Expand allocations inside of maps when possible.
    module Futhark.Pass.ExpandAllocations (expandAllocations) where

    import Control.Monad
    import Control.Monad.Except
    import Control.Monad.Reader
    import Control.Monad.State
    import Control.Monad.Writer
    import Data.Bifunctor
    import Data.Either (rights)
    import Data.List (find, foldl')
    import Data.Map.Strict qualified as M
    import Data.Maybe
    import Data.Sequence qualified as Seq
    import Futhark.Analysis.Alias as Alias
    import Futhark.Analysis.SymbolTable qualified as ST
    import Futhark.Error
    import Futhark.IR
    import Futhark.IR.GPU.Simplify qualified as GPU
    import Futhark.IR.GPUMem
    import Futhark.IR.Mem.LMAD qualified as LMAD
    import Futhark.MonadFreshNames
    import Futhark.Optimise.Simplify.Rep (addScopeWisdom)
    import Futhark.Pass
    import Futhark.Pass.ExplicitAllocations.GPU (explicitAllocationsInStms)
    import Futhark.Pass.ExtractKernels.BlockedKernel (nonSegRed)
    import Futhark.Pass.ExtractKernels.ToGPU (segThread)
    import Futhark.Tools
    import Futhark.Transform.CopyPropagate (copyPropagateInFun)
    import Futhark.Transform.Rename (renameStm)
    import Futhark.Transform.Substitute
    import Futhark.Util (mapAccumLM)
    import Prelude hiding (quot)

    -- | The memory expansion pass definition.
    expandAllocations :: Pass GPUMem GPUMem
    expandAllocations :: Pass GPUMem GPUMem
    expandAllocations =
      String
    -> String
    -> (Prog GPUMem -> PassM (Prog GPUMem))
    -> Pass GPUMem GPUMem
    forall fromrep torep.
    String
    -> String
    -> (Prog fromrep -> PassM (Prog torep))
    -> Pass fromrep torep
    Pass String
    "expand allocations" String
    "Expand allocations" ((Prog GPUMem -> PassM (Prog GPUMem)) -> Pass GPUMem GPUMem)
    -> (Prog GPUMem -> PassM (Prog GPUMem)) -> Pass GPUMem GPUMem
    forall a b. (a -> b) -> a -> b
    $
        \Prog GPUMem
    prog -> do
          consts' <-
            (VNameSource -> (Stms GPUMem, VNameSource)) -> PassM (Stms GPUMem)
    forall (m :: * -> *) a.
    MonadFreshNames m =>
    (VNameSource -> (a, VNameSource)) -> m a
    modifyNameSource ((VNameSource -> (Stms GPUMem, VNameSource))
     -> PassM (Stms GPUMem))
    -> (VNameSource -> (Stms GPUMem, VNameSource))
    -> PassM (Stms GPUMem)
    forall a b. (a -> b) -> a -> b
    $
              Either String (Stms GPUMem, VNameSource)
    -> (Stms GPUMem, VNameSource)
    forall a. Either String a -> a
    limitationOnLeft
                (Either String (Stms GPUMem, VNameSource)
     -> (Stms GPUMem, VNameSource))
    -> (VNameSource -> Either String (Stms GPUMem, VNameSource))
    -> VNameSource
    -> (Stms GPUMem, VNameSource)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . StateT VNameSource (Either String) (Stms GPUMem)
    -> VNameSource -> Either String (Stms GPUMem, VNameSource)
    forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
    runStateT (ReaderT
      (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    -> Scope GPUMem -> StateT VNameSource (Either String) (Stms GPUMem)
    forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
    runReaderT (Stms GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    transformStms (Prog GPUMem -> Stms GPUMem
    forall rep. Prog rep -> Stms rep
    progConsts Prog GPUMem
    prog)) Scope GPUMem
    forall a. Monoid a => a
    mempty)
          funs' <- mapM (transformFunDef $ scopeOf consts') (progFuns prog)
          pure $ prog {progConsts = consts', progFuns = funs'}

    -- Cannot use intraproceduralTransformation because it might create
    -- duplicate size keys (which are not fixed by renamer, and size
    -- keys must currently be globally unique).

    type ExpandM = ReaderT (Scope GPUMem) (StateT VNameSource (Either String))

    limitationOnLeft :: Either String a -> a
    limitationOnLeft :: forall a. Either String a -> a
    limitationOnLeft = (String -> a) -> (a -> a) -> Either String a -> a
    forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
    either String -> a
    forall a. String -> a
    compilerLimitationS a -> a
    forall a. a -> a
    id

    transformFunDef ::
      Scope GPUMem ->
      FunDef GPUMem ->
      PassM (FunDef GPUMem)
    transformFunDef :: Scope GPUMem -> FunDef GPUMem -> PassM (FunDef GPUMem)
    transformFunDef Scope GPUMem
    scope FunDef GPUMem
    fundec = do
      body' <- (VNameSource -> (Body GPUMem, VNameSource)) -> PassM (Body GPUMem)
    forall (m :: * -> *) a.
    MonadFreshNames m =>
    (VNameSource -> (a, VNameSource)) -> m a
    modifyNameSource ((VNameSource -> (Body GPUMem, VNameSource))
     -> PassM (Body GPUMem))
    -> (VNameSource -> (Body GPUMem, VNameSource))
    -> PassM (Body GPUMem)
    forall a b. (a -> b) -> a -> b
    $ Either String (Body GPUMem, VNameSource)
    -> (Body GPUMem, VNameSource)
    forall a. Either String a -> a
    limitationOnLeft (Either String (Body GPUMem, VNameSource)
     -> (Body GPUMem, VNameSource))
    -> (VNameSource -> Either String (Body GPUMem, VNameSource))
    -> VNameSource
    -> (Body GPUMem, VNameSource)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . StateT VNameSource (Either String) (Body GPUMem)
    -> VNameSource -> Either String (Body GPUMem, VNameSource)
    forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
    runStateT (ReaderT
      (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    -> Scope GPUMem -> StateT VNameSource (Either String) (Body GPUMem)
    forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
    runReaderT ReaderT
      (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    m Scope GPUMem
    forall a. Monoid a => a
    mempty)
      copyPropagateInFun
        simpleGPUMem
        (ST.fromScope (addScopeWisdom scope))
        fundec {funDefBody = body'}
      where
        m :: ReaderT
      (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    m =
          Scope GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    forall a.
    Scope GPUMem
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope Scope GPUMem
    scope (ReaderT
       (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
     -> ReaderT
          (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem))
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    forall a b. (a -> b) -> a -> b
    $
            FunDef GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    forall rep a (m :: * -> *) b.
    (Scoped rep a, LocalScope rep m) =>
    a -> m b -> m b
    inScopeOf FunDef GPUMem
    fundec (ReaderT
       (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
     -> ReaderT
          (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem))
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    forall a b. (a -> b) -> a -> b
    $
              Body GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    transformBody (Body GPUMem
     -> ReaderT
          (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem))
    -> Body GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    forall a b. (a -> b) -> a -> b
    $
                FunDef GPUMem -> Body GPUMem
    forall rep. FunDef rep -> Body rep
    funDefBody FunDef GPUMem
    fundec

    transformBody :: Body GPUMem -> ExpandM (Body GPUMem)
    transformBody :: Body GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    transformBody (Body () Stms GPUMem
    stms Result
    res) = BodyDec GPUMem -> Stms GPUMem -> Result -> Body GPUMem
    forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
    Body () (Stms GPUMem -> Result -> Body GPUMem)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         (Result -> Body GPUMem)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Stms GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    transformStms Stms GPUMem
    stms ReaderT
      (Scope GPUMem)
      (StateT VNameSource (Either String))
      (Result -> Body GPUMem)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) Result
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    forall a b.
    ReaderT
      (Scope GPUMem) (StateT VNameSource (Either String)) (a -> b)
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Result
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) Result
    forall a.
    a -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Result
    res

    transformLambda :: Lambda GPUMem -> ExpandM (Lambda GPUMem)
    transformLambda :: Lambda GPUMem -> ExpandM (Lambda GPUMem)
    transformLambda (Lambda [LParam GPUMem]
    params [TypeBase (ShapeBase SubExp) NoUniqueness]
    ret Body GPUMem
    body) =
      [LParam GPUMem]
    -> [TypeBase (ShapeBase SubExp) NoUniqueness]
    -> Body GPUMem
    -> Lambda GPUMem
    forall rep.
    [LParam rep]
    -> [TypeBase (ShapeBase SubExp) NoUniqueness]
    -> Body rep
    -> Lambda rep
    Lambda [LParam GPUMem]
    params [TypeBase (ShapeBase SubExp) NoUniqueness]
    ret
        (Body GPUMem -> Lambda GPUMem)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    -> ExpandM (Lambda GPUMem)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Scope GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    forall a.
    Scope GPUMem
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope ([LParam GPUMem] -> Scope GPUMem
    forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfLParams [LParam GPUMem]
    params) (Body GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    transformBody Body GPUMem
    body)

    transformStms :: Stms GPUMem -> ExpandM (Stms GPUMem)
    transformStms :: Stms GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    transformStms Stms GPUMem
    stms =
      Stms GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    forall rep a (m :: * -> *) b.
    (Scoped rep a, LocalScope rep m) =>
    a -> m b -> m b
    inScopeOf Stms GPUMem
    stms (ReaderT
       (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
     -> ReaderT
          (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem))
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    forall a b. (a -> b) -> a -> b
    $ [Stms GPUMem] -> Stms GPUMem
    forall a. Monoid a => [a] -> a
    mconcat ([Stms GPUMem] -> Stms GPUMem)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) [Stms GPUMem]
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Stm GPUMem
     -> ReaderT
          (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem))
    -> [Stm GPUMem]
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) [Stms GPUMem]
    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 Stm GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    transformStm (Stms GPUMem -> [Stm GPUMem]
    forall rep. Stms rep -> [Stm rep]
    stmsToList Stms GPUMem
    stms)

    transformStm :: Stm GPUMem -> ExpandM (Stms GPUMem)
    -- It is possible that we are unable to expand allocations in some
    -- code versions.  If so, we can remove the offending branch.  Only if
    -- all versions fail do we propagate the error.
    -- FIXME: this can remove safety checks if the default branch fails!
    transformStm :: Stm GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    transformStm (Let Pat (LetDec GPUMem)
    pat StmAux (ExpDec GPUMem)
    aux (Match [SubExp]
    cond [Case (Body GPUMem)]
    cases Body GPUMem
    defbody (MatchDec [BranchType GPUMem]
    ts MatchSort
    MatchEquiv))) = do
      let onCase :: Case (Body GPUMem)
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         (Either String (Case (Body GPUMem)))
    onCase (Case [Maybe PrimValue]
    vs Body GPUMem
    body) =
            (Case (Body GPUMem) -> Either String (Case (Body GPUMem))
    forall a b. b -> Either a b
    Right (Case (Body GPUMem) -> Either String (Case (Body GPUMem)))
    -> (Body GPUMem -> Case (Body GPUMem))
    -> Body GPUMem
    -> Either String (Case (Body GPUMem))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [Maybe PrimValue] -> Body GPUMem -> Case (Body GPUMem)
    forall body. [Maybe PrimValue] -> body -> Case body
    Case [Maybe PrimValue]
    vs (Body GPUMem -> Either String (Case (Body GPUMem)))
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         (Either String (Case (Body GPUMem)))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Body GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    transformBody Body GPUMem
    body) ReaderT
      (Scope GPUMem)
      (StateT VNameSource (Either String))
      (Either String (Case (Body GPUMem)))
    -> (String
        -> ReaderT
             (Scope GPUMem)
             (StateT VNameSource (Either String))
             (Either String (Case (Body GPUMem))))
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         (Either String (Case (Body GPUMem)))
    forall a.
    ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    -> (String
        -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a)
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall e (m :: * -> *) a.
    MonadError e m =>
    m a -> (e -> m a) -> m a
    `catchError` (Either String (Case (Body GPUMem))
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         (Either String (Case (Body GPUMem)))
    forall a.
    a -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Either String (Case (Body GPUMem))
     -> ReaderT
          (Scope GPUMem)
          (StateT VNameSource (Either String))
          (Either String (Case (Body GPUMem))))
    -> (String -> Either String (Case (Body GPUMem)))
    -> String
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         (Either String (Case (Body GPUMem)))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . String -> Either String (Case (Body GPUMem))
    forall a b. a -> Either a b
    Left)
      cases' <- [Either String (Case (Body GPUMem))] -> [Case (Body GPUMem)]
    forall a b. [Either a b] -> [b]
    rights ([Either String (Case (Body GPUMem))] -> [Case (Body GPUMem)])
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         [Either String (Case (Body GPUMem))]
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         [Case (Body GPUMem)]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Case (Body GPUMem)
     -> ReaderT
          (Scope GPUMem)
          (StateT VNameSource (Either String))
          (Either String (Case (Body GPUMem))))
    -> [Case (Body GPUMem)]
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         [Either String (Case (Body GPUMem))]
    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 Case (Body GPUMem)
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         (Either String (Case (Body GPUMem)))
    onCase [Case (Body GPUMem)]
    cases
      defbody' <- (Right <$> transformBody defbody) `catchError` (pure . Left)
      case (cases', defbody') of
        ([], Left String
    e) ->
          String
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    forall a.
    String
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall e (m :: * -> *) a. MonadError e m => e -> m a
    throwError String
    e
        (Case (Body GPUMem)
    _ : [Case (Body GPUMem)]
    _, Left String
    _) ->
          Stms GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    forall a.
    a -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Stms GPUMem
     -> ReaderT
          (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem))
    -> Stms GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    forall a b. (a -> b) -> a -> b
    $ Stm GPUMem -> Stms GPUMem
    forall rep. Stm rep -> Stms rep
    oneStm (Stm GPUMem -> Stms GPUMem) -> Stm GPUMem -> Stms GPUMem
    forall a b. (a -> b) -> a -> b
    $ Pat (LetDec GPUMem)
    -> StmAux (ExpDec GPUMem) -> Exp GPUMem -> Stm GPUMem
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let Pat (LetDec GPUMem)
    pat StmAux (ExpDec GPUMem)
    aux (Exp GPUMem -> Stm GPUMem) -> Exp GPUMem -> Stm GPUMem
    forall a b. (a -> b) -> a -> b
    $ [SubExp]
    -> [Case (Body GPUMem)]
    -> Body GPUMem
    -> MatchDec (BranchType GPUMem)
    -> Exp GPUMem
    forall rep.
    [SubExp]
    -> [Case (Body rep)]
    -> Body rep
    -> MatchDec (BranchType rep)
    -> Exp rep
    Match [SubExp]
    cond ([Case (Body GPUMem)] -> [Case (Body GPUMem)]
    forall a. HasCallStack => [a] -> [a]
    init [Case (Body GPUMem)]
    cases') (Case (Body GPUMem) -> Body GPUMem
    forall body. Case body -> body
    caseBody (Case (Body GPUMem) -> Body GPUMem)
    -> Case (Body GPUMem) -> Body GPUMem
    forall a b. (a -> b) -> a -> b
    $ [Case (Body GPUMem)] -> Case (Body GPUMem)
    forall a. HasCallStack => [a] -> a
    last [Case (Body GPUMem)]
    cases') ([BranchTypeMem] -> MatchSort -> MatchDec BranchTypeMem
    forall rt. [rt] -> MatchSort -> MatchDec rt
    MatchDec [BranchType GPUMem]
    [BranchTypeMem]
    ts MatchSort
    MatchEquiv)
        ([Case (Body GPUMem)]
    _, Right Body GPUMem
    defbody'') ->
          Stms GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    forall a.
    a -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Stms GPUMem
     -> ReaderT
          (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem))
    -> Stms GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPUMem)
    forall a b. (a -> b) -> a -> b
    $ Stm GPUMem -> Stms GPUMem
    forall rep. Stm rep -> Stms rep
    oneStm (Stm GPUMem -> Stms GPUMem) -> Stm GPUMem -> Stms GPUMem
    forall a b. (a -> b) -> a -> b
    $ Pat (LetDec GPUMem)
    -> StmAux (ExpDec GPUMem) -> Exp GPUMem -> Stm GPUMem
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let Pat (LetDec GPUMem)
    pat StmAux (ExpDec GPUMem)
    aux (Exp GPUMem -> Stm GPUMem) -> Exp GPUMem -> Stm GPUMem
    forall a b. (a -> b) -> a -> b
    $ [SubExp]
    -> [Case (Body GPUMem)]
    -> Body GPUMem
    -> MatchDec (BranchType GPUMem)
    -> Exp GPUMem
    forall rep.
    [SubExp]
    -> [Case (Body rep)]
    -> Body rep
    -> MatchDec (BranchType rep)
    -> Exp rep
    Match [SubExp]
    cond [Case (Body GPUMem)]
    cases' Body GPUMem
    defbody'' ([BranchTypeMem] -> MatchSort -> MatchDec BranchTypeMem
    forall rt. [rt] -> MatchSort -> MatchDec rt
    MatchDec [BranchType GPUMem]
    [BranchTypeMem]
    ts MatchSort
    MatchEquiv)
    transformStm (Let Pat (LetDec GPUMem)
    pat StmAux (ExpDec GPUMem)
    aux Exp GPUMem
    e) = do
      (stms, e') <- Exp GPUMem -> ExpandM (Stms GPUMem, Exp GPUMem)
    transformExp (Exp GPUMem -> ExpandM (Stms GPUMem, Exp GPUMem))
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Exp GPUMem)
    -> ExpandM (Stms GPUMem, Exp GPUMem)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Mapper
      GPUMem
      GPUMem
      (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
    -> Exp GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Exp GPUMem)
    forall (m :: * -> *) frep trep.
    Monad m =>
    Mapper frep trep m -> Exp frep -> m (Exp trep)
    mapExpM Mapper
      GPUMem
      GPUMem
      (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
    transform Exp GPUMem
    e
      pure $ stms <> oneStm (Let pat aux e')
      where
        transform :: Mapper
      GPUMem
      GPUMem
      (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
    transform =
          Mapper
      GPUMem
      GPUMem
      (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
    forall rep (m :: * -> *). Monad m => Mapper rep rep m
    identityMapper
            { mapOnBody = \Scope GPUMem
    scope -> Scope GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    forall a.
    Scope GPUMem
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope Scope GPUMem
    scope (ReaderT
       (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
     -> ReaderT
          (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem))
    -> (Body GPUMem
        -> ReaderT
             (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem))
    -> Body GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Body GPUMem
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Body GPUMem)
    transformBody
            }

    transformExp :: Exp GPUMem -> ExpandM (Stms GPUMem, Exp GPUMem)
    transformExp :: Exp GPUMem -> ExpandM (Stms GPUMem, Exp GPUMem)
    transformExp (Op (Inner (SegOp (SegMap SegLevel
    lvl SegSpace
    space [TypeBase (ShapeBase SubExp) NoUniqueness]
    ts KernelBody GPUMem
    kbody)))) = do
      (alloc_stms, (lvl', _, kbody')) <- SegLevel
    -> SegSpace
    -> [Lambda GPUMem]
    -> KernelBody GPUMem
    -> ExpandM
         (Stms GPUMem, (SegLevel, [Lambda GPUMem], KernelBody GPUMem))
    transformScanRed SegLevel
    lvl SegSpace
    space [] KernelBody GPUMem
    kbody
      pure
        ( alloc_stms,
          Op $ Inner $ SegOp $ SegMap lvl' space ts kbody'
        )
    transformExp (Op (Inner (SegOp (SegRed SegLevel
    lvl SegSpace
    space [TypeBase (ShapeBase SubExp) NoUniqueness]
    ts KernelBody GPUMem
    kbody [SegBinOp GPUMem]
    reds)))) = do
      (alloc_stms, (lvl', lams, kbody')) <-
        SegLevel
    -> SegSpace
    -> [Lambda GPUMem]
    -> KernelBody GPUMem
    -> ExpandM
         (Stms GPUMem, (SegLevel, [Lambda GPUMem], KernelBody GPUMem))
    transformScanRed SegLevel
    lvl SegSpace
    space ((SegBinOp GPUMem -> Lambda GPUMem)
    -> [SegBinOp GPUMem] -> [Lambda GPUMem]
    forall a b. (a -> b) -> [a] -> [b]
    map SegBinOp GPUMem -> Lambda GPUMem
    forall rep. SegBinOp rep -> Lambda rep
    segBinOpLambda [SegBinOp GPUMem]
    reds) KernelBody GPUMem
    kbody
      let reds' = (SegBinOp GPUMem -> Lambda GPUMem -> SegBinOp GPUMem)
    -> [SegBinOp GPUMem] -> [Lambda GPUMem] -> [SegBinOp GPUMem]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith (\SegBinOp GPUMem
    red Lambda GPUMem
    lam -> SegBinOp GPUMem
    red {segBinOpLambda = lam}) [SegBinOp GPUMem]
    reds [Lambda GPUMem]
    lams
      pure
        ( alloc_stms,
          Op $ Inner $ SegOp $ SegRed lvl' space ts kbody' reds'
        )
    transformExp (Op (Inner (SegOp (SegScan SegLevel
    lvl SegSpace
    space [TypeBase (ShapeBase SubExp) NoUniqueness]
    ts KernelBody GPUMem
    kbody [SegBinOp GPUMem]
    scans)))) = do
      (alloc_stms, (lvl', lams, kbody')) <-
        SegLevel
    -> SegSpace
    -> [Lambda GPUMem]
    -> KernelBody GPUMem
    -> ExpandM
         (Stms GPUMem, (SegLevel, [Lambda GPUMem], KernelBody GPUMem))
    transformScanRed SegLevel
    lvl SegSpace
    space ((SegBinOp GPUMem -> Lambda GPUMem)
    -> [SegBinOp GPUMem] -> [Lambda GPUMem]
    forall a b. (a -> b) -> [a] -> [b]
    map SegBinOp GPUMem -> Lambda GPUMem
    forall rep. SegBinOp rep -> Lambda rep
    segBinOpLambda [SegBinOp GPUMem]
    scans) KernelBody GPUMem
    kbody
      let scans' = (SegBinOp GPUMem -> Lambda GPUMem -> SegBinOp GPUMem)
    -> [SegBinOp GPUMem] -> [Lambda GPUMem] -> [SegBinOp GPUMem]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith (\SegBinOp GPUMem
    red Lambda GPUMem
    lam -> SegBinOp GPUMem
    red {segBinOpLambda = lam}) [SegBinOp GPUMem]
    scans [Lambda GPUMem]
    lams
      pure
        ( alloc_stms,
          Op $ Inner $ SegOp $ SegScan lvl' space ts kbody' scans'
        )
    transformExp (Op (Inner (SegOp (SegHist SegLevel
    lvl SegSpace
    space [TypeBase (ShapeBase SubExp) NoUniqueness]
    ts KernelBody GPUMem
    kbody [HistOp GPUMem]
    ops)))) = do
      (alloc_stms, (lvl', lams', kbody')) <- SegLevel
    -> SegSpace
    -> [Lambda GPUMem]
    -> KernelBody GPUMem
    -> ExpandM
         (Stms GPUMem, (SegLevel, [Lambda GPUMem], KernelBody GPUMem))
    transformScanRed SegLevel
    lvl SegSpace
    space [Lambda GPUMem]
    lams KernelBody GPUMem
    kbody
      let ops' = (HistOp GPUMem -> Lambda GPUMem -> HistOp GPUMem)
    -> [HistOp GPUMem] -> [Lambda GPUMem] -> [HistOp GPUMem]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith HistOp GPUMem -> Lambda GPUMem -> HistOp GPUMem
    forall {rep} {rep}. HistOp rep -> Lambda rep -> HistOp rep
    onOp [HistOp GPUMem]
    ops [Lambda GPUMem]
    lams'
      pure
        ( alloc_stms,
          Op $ Inner $ SegOp $ SegHist lvl' space ts kbody' ops'
        )
      where
        lams :: [Lambda GPUMem]
    lams = (HistOp GPUMem -> Lambda GPUMem)
    -> [HistOp GPUMem] -> [Lambda GPUMem]
    forall a b. (a -> b) -> [a] -> [b]
    map HistOp GPUMem -> Lambda GPUMem
    forall rep. HistOp rep -> Lambda rep
    histOp [HistOp GPUMem]
    ops
        onOp :: HistOp rep -> Lambda rep -> HistOp rep
    onOp HistOp rep
    op Lambda rep
    lam = HistOp rep
    op {histOp = lam}
    transformExp (WithAcc [WithAccInput GPUMem]
    inputs Lambda GPUMem
    lam) = do
      lam' <- Lambda GPUMem -> ExpandM (Lambda GPUMem)
    transformLambda Lambda GPUMem
    lam
      (input_alloc_stms, inputs') <- mapAndUnzipM onInput inputs
      pure
        ( mconcat input_alloc_stms,
          WithAcc inputs' lam'
        )
      where
        onInput :: (ShapeBase SubExp, b, Maybe (Lambda GPUMem, b))
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         (Stms GPUMem, (ShapeBase SubExp, b, Maybe (Lambda GPUMem, b)))
    onInput (ShapeBase SubExp
    shape, b
    arrs, Maybe (Lambda GPUMem, b)
    Nothing) =
          (Stms GPUMem, (ShapeBase SubExp, b, Maybe (Lambda GPUMem, b)))
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         (Stms GPUMem, (ShapeBase SubExp, b, Maybe (Lambda GPUMem, b)))
    forall a.
    a -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Stms GPUMem
    forall a. Monoid a => a
    mempty, (ShapeBase SubExp
    shape, b
    arrs, Maybe (Lambda GPUMem, b)
    forall a. Maybe a
    Nothing))
        onInput (ShapeBase SubExp
    shape, b
    arrs, Just (Lambda GPUMem
    op_lam, b
    nes)) = do
          bound_outside <- (Scope GPUMem -> Names)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) Names
    forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
    asks ((Scope GPUMem -> Names)
     -> ReaderT
          (Scope GPUMem) (StateT VNameSource (Either String)) Names)
    -> (Scope GPUMem -> Names)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) Names
    forall a b. (a -> b) -> a -> b
    $ [VName] -> Names
    namesFromList ([VName] -> Names)
    -> (Scope GPUMem -> [VName]) -> Scope GPUMem -> Names
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Scope GPUMem -> [VName]
    forall k a. Map k a -> [k]
    M.keys
          let -- XXX: fake a SegLevel, which we don't have here.  We will not
              -- use it for anything, as we will not allow irregular
              -- allocations inside the update function.
              lvl = SegVirt -> Maybe KernelGrid -> SegLevel
    SegThread SegVirt
    SegNoVirt Maybe KernelGrid
    forall a. Maybe a
    Nothing
              (op_lam', lam_allocs) =
                extractLambdaAllocations (lvl, [0]) bound_outside mempty op_lam
              variantAlloc ((SegLevel, [TPrimExp Int64 VName])
    _, Var VName
    v, Space
    _) = VName
    v VName -> Names -> Bool
    `notNameIn` Names
    bound_outside
              variantAlloc ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space)
    _ = Bool
    False
              (variant_allocs, invariant_allocs) = M.partition variantAlloc lam_allocs

          case M.elems variant_allocs of
            ((SegLevel, [TPrimExp Int64 VName])
    _, SubExp
    v, Space
    _) : [((SegLevel, [TPrimExp Int64 VName]), SubExp, Space)]
    _ ->
              String
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) ()
    forall a.
    String
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall e (m :: * -> *) a. MonadError e m => e -> m a
    throwError (String
     -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) ())
    -> String
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) ()
    forall a b. (a -> b) -> a -> b
    $
                String
    "Cannot handle un-sliceable allocation size: "
                  String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ SubExp -> String
    forall a. Pretty a => a -> String
    prettyString SubExp
    v
                  String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    "\nLikely cause: irregular nested operations inside accumulator update operator."
            [] ->
              ()
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) ()
    forall a.
    a -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

          let num_is = ShapeBase SubExp -> Int
    forall a. ArrayShape a => a -> Int
    shapeRank ShapeBase SubExp
    shape
              is = (Param (MemInfo SubExp NoUniqueness MemBind) -> VName)
    -> [Param (MemInfo SubExp NoUniqueness MemBind)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (MemInfo SubExp NoUniqueness MemBind) -> VName
    forall dec. Param dec -> VName
    paramName ([Param (MemInfo SubExp NoUniqueness MemBind)] -> [VName])
    -> [Param (MemInfo SubExp NoUniqueness MemBind)] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Int
    -> [Param (MemInfo SubExp NoUniqueness MemBind)]
    -> [Param (MemInfo SubExp NoUniqueness MemBind)]
    forall a. Int -> [a] -> [a]
    take Int
    num_is ([Param (MemInfo SubExp NoUniqueness MemBind)]
     -> [Param (MemInfo SubExp NoUniqueness MemBind)])
    -> [Param (MemInfo SubExp NoUniqueness MemBind)]
    -> [Param (MemInfo SubExp NoUniqueness MemBind)]
    forall a b. (a -> b) -> a -> b
    $ Lambda GPUMem -> [LParam GPUMem]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda GPUMem
    op_lam
          (alloc_stms, alloc_offsets) <-
            genericExpandedInvariantAllocations (const $ const (shape, map le64 is)) invariant_allocs

          scope <- askScope
          let scope' = Lambda GPUMem -> Scope GPUMem
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf Lambda GPUMem
    op_lam Scope GPUMem -> Scope GPUMem -> Scope GPUMem
    forall a. Semigroup a => a -> a -> a
    <> Scope GPUMem
    scope Scope GPUMem -> Scope GPUMem -> Scope GPUMem
    forall a. Semigroup a => a -> a -> a
    <> Stms GPUMem -> Scope GPUMem
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf Stms GPUMem
    alloc_stms
          either throwError pure <=< runOffsetM scope' $ do
            op_lam'' <- offsetMemoryInLambda alloc_offsets op_lam'
            pure (alloc_stms, (shape, arrs, Just (op_lam'', nes)))
    transformExp Exp GPUMem
    e =
      (Stms GPUMem, Exp GPUMem) -> ExpandM (Stms GPUMem, Exp GPUMem)
    forall a.
    a -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Stms GPUMem
    forall a. Monoid a => a
    mempty, Exp GPUMem
    e)

    ensureGridKnown :: SegLevel -> ExpandM (Stms GPUMem, SegLevel, KernelGrid)
    ensureGridKnown :: SegLevel -> ExpandM (Stms GPUMem, SegLevel, KernelGrid)
    ensureGridKnown SegLevel
    lvl =
      case SegLevel
    lvl of
        SegThread SegVirt
    _ (Just KernelGrid
    grid) -> (Stms GPUMem, SegLevel, KernelGrid)
    -> ExpandM (Stms GPUMem, SegLevel, KernelGrid)
    forall a.
    a -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Stms GPUMem
    forall a. Monoid a => a
    mempty, SegLevel
    lvl, KernelGrid
    grid)
        SegBlock SegVirt
    _ (Just KernelGrid
    grid) -> (Stms GPUMem, SegLevel, KernelGrid)
    -> ExpandM (Stms GPUMem, SegLevel, KernelGrid)
    forall a.
    a -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Stms GPUMem
    forall a. Monoid a => a
    mempty, SegLevel
    lvl, KernelGrid
    grid)
        SegThread SegVirt
    virt Maybe KernelGrid
    Nothing -> (Maybe KernelGrid -> SegLevel)
    -> ExpandM (Stms GPUMem, SegLevel, KernelGrid)
    forall {somerep} {rep} {op :: * -> *} {m :: * -> *} {b}.
    (FParamInfo somerep ~ FParamInfo rep, OpC rep ~ MemOp (HostOp op),
     LetDec somerep ~ LetDec rep, LParamInfo somerep ~ LParamInfo rep,
     HasScope somerep m, MonadFreshNames m, BuilderOps rep, IsOp op,
     RephraseOp op, Pretty (op rep), Rename (op rep),
     Substitute (op rep), FreeIn (op rep), Show (op rep),
     Ord (op rep)) =>
    (Maybe KernelGrid -> b) -> m (Stms rep, b, KernelGrid)
    mkGrid (SegVirt -> Maybe KernelGrid -> SegLevel
    SegThread SegVirt
    virt)
        SegBlock SegVirt
    virt Maybe KernelGrid
    Nothing -> (Maybe KernelGrid -> SegLevel)
    -> ExpandM (Stms GPUMem, SegLevel, KernelGrid)
    forall {somerep} {rep} {op :: * -> *} {m :: * -> *} {b}.
    (FParamInfo somerep ~ FParamInfo rep, OpC rep ~ MemOp (HostOp op),
     LetDec somerep ~ LetDec rep, LParamInfo somerep ~ LParamInfo rep,
     HasScope somerep m, MonadFreshNames m, BuilderOps rep, IsOp op,
     RephraseOp op, Pretty (op rep), Rename (op rep),
     Substitute (op rep), FreeIn (op rep), Show (op rep),
     Ord (op rep)) =>
    (Maybe KernelGrid -> b) -> m (Stms rep, b, KernelGrid)
    mkGrid (SegVirt -> Maybe KernelGrid -> SegLevel
    SegBlock SegVirt
    virt)
        SegThreadInBlock {} -> String -> ExpandM (Stms GPUMem, SegLevel, KernelGrid)
    forall a. HasCallStack => String -> a
    error String
    "ensureGridKnown: SegThreadInBlock"
      where
        mkGrid :: (Maybe KernelGrid -> b) -> m (Stms rep, b, KernelGrid)
    mkGrid Maybe KernelGrid -> b
    f = do
          (grid, stms) <-
            Builder rep KernelGrid -> m (KernelGrid, Stms rep)
    forall (m :: * -> *) somerep rep a.
    (MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
    Builder rep a -> m (a, Stms rep)
    runBuilder (Builder rep KernelGrid -> m (KernelGrid, Stms rep))
    -> Builder rep KernelGrid -> m (KernelGrid, Stms rep)
    forall a b. (a -> b) -> a -> b
    $
              Count NumBlocks SubExp -> Count BlockSize SubExp -> KernelGrid
    KernelGrid
                (Count NumBlocks SubExp -> Count BlockSize SubExp -> KernelGrid)
    -> BuilderT rep (State VNameSource) (Count NumBlocks SubExp)
    -> BuilderT
         rep (State VNameSource) (Count BlockSize SubExp -> KernelGrid)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (SubExp -> Count NumBlocks SubExp
    forall {k} (u :: k) e. e -> Count u e
    Count (SubExp -> Count NumBlocks SubExp)
    -> BuilderT rep (State VNameSource) SubExp
    -> BuilderT rep (State VNameSource) (Count NumBlocks SubExp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> String -> SizeClass -> BuilderT rep (State VNameSource) SubExp
    forall {m :: * -> *} {op :: * -> *}.
    (OpC (Rep m) ~ MemOp (HostOp op), MonadBuilder m,
     Pretty (op (Rep m)), Rename (op (Rep m)), Substitute (op (Rep m)),
     FreeIn (op (Rep m)), IsOp op, RephraseOp op, Show (op (Rep m)),
     Ord (op (Rep m))) =>
    String -> SizeClass -> m SubExp
    getSize String
    "num_tblocks" SizeClass
    SizeGrid)
                BuilderT
      rep (State VNameSource) (Count BlockSize SubExp -> KernelGrid)
    -> BuilderT rep (State VNameSource) (Count BlockSize SubExp)
    -> Builder rep KernelGrid
    forall a b.
    BuilderT rep (State VNameSource) (a -> b)
    -> BuilderT rep (State VNameSource) a
    -> BuilderT rep (State VNameSource) b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (SubExp -> Count BlockSize SubExp
    forall {k} (u :: k) e. e -> Count u e
    Count (SubExp -> Count BlockSize SubExp)
    -> BuilderT rep (State VNameSource) SubExp
    -> BuilderT rep (State VNameSource) (Count BlockSize SubExp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> String -> SizeClass -> BuilderT rep (State VNameSource) SubExp
    forall {m :: * -> *} {op :: * -> *}.
    (OpC (Rep m) ~ MemOp (HostOp op), MonadBuilder m,
     Pretty (op (Rep m)), Rename (op (Rep m)), Substitute (op (Rep m)),
     FreeIn (op (Rep m)), IsOp op, RephraseOp op, Show (op (Rep m)),
     Ord (op (Rep m))) =>
    String -> SizeClass -> m SubExp
    getSize String
    "tblock_size" SizeClass
    SizeThreadBlock)
          pure (stms, f $ Just grid, grid)

        getSize :: String -> SizeClass -> m SubExp
    getSize String
    desc SizeClass
    size_class = do
          size_key <- String -> Name
    nameFromString (String -> Name) -> (VName -> String) -> VName -> Name
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> String
    forall a. Pretty a => a -> String
    prettyString (VName -> Name) -> m VName -> m Name
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> String -> m VName
    forall (m :: * -> *). MonadFreshNames m => String -> m VName
    newVName String
    desc
          letSubExp desc $ Op $ Inner $ SizeOp $ GetSize size_key size_class

    transformScanRed ::
      SegLevel ->
      SegSpace ->
      [Lambda GPUMem] ->
      KernelBody GPUMem ->
      ExpandM (Stms GPUMem, (SegLevel, [Lambda GPUMem], KernelBody GPUMem))
    transformScanRed :: SegLevel
    -> SegSpace
    -> [Lambda GPUMem]
    -> KernelBody GPUMem
    -> ExpandM
         (Stms GPUMem, (SegLevel, [Lambda GPUMem], KernelBody GPUMem))
    transformScanRed SegLevel
    lvl SegSpace
    space [Lambda GPUMem]
    ops KernelBody GPUMem
    kbody = do
      bound_outside <- (Scope GPUMem -> Names)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) Names
    forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
    asks ((Scope GPUMem -> Names)
     -> ReaderT
          (Scope GPUMem) (StateT VNameSource (Either String)) Names)
    -> (Scope GPUMem -> Names)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) Names
    forall a b. (a -> b) -> a -> b
    $ [VName] -> Names
    namesFromList ([VName] -> Names)
    -> (Scope GPUMem -> [VName]) -> Scope GPUMem -> Names
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Scope GPUMem -> [VName]
    forall k a. Map k a -> [k]
    M.keys
      let user = (SegLevel
    lvl, [VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 (VName -> TPrimExp Int64 VName) -> VName -> TPrimExp Int64 VName
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> VName
    segFlat SegSpace
    space])
          (kbody', kbody_allocs) =
            extractKernelBodyAllocations user bound_outside bound_in_kernel kbody
          (ops', ops_allocs) =
            unzip $ map (extractLambdaAllocations user bound_outside mempty) ops
          variantAlloc ((SegLevel, [TPrimExp Int64 VName])
    _, Var VName
    v, Space
    _) = VName
    v VName -> Names -> Bool
    `notNameIn` Names
    bound_outside
          variantAlloc ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space)
    _ = Bool
    False
          (variant_allocs, invariant_allocs) =
            M.partition variantAlloc $ kbody_allocs <> mconcat ops_allocs
          badVariant ((SegLevel, [TPrimExp Int64 VName])
    _, Var VName
    v, Space
    _) = VName
    v VName -> Names -> Bool
    `notNameIn` Names
    bound_in_kernel
          badVariant ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space)
    _ = Bool
    False

      case find badVariant $ M.elems variant_allocs of
        Just ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space)
    v ->
          String
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) ()
    forall a.
    String
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall e (m :: * -> *) a. MonadError e m => e -> m a
    throwError (String
     -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) ())
    -> String
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) ()
    forall a b. (a -> b) -> a -> b
    $
            String
    "Cannot handle un-sliceable allocation size: "
              String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space) -> String
    forall a. Pretty a => a -> String
    prettyString ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space)
    v
              String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    "\nLikely cause: irregular nested operations inside parallel constructs."
        Maybe ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space)
    Nothing ->
          ()
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) ()
    forall a.
    a -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

      case lvl of
        SegBlock {}
          | Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ Extraction -> Bool
    forall a. Map VName a -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null Extraction
    variant_allocs ->
              String
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) ()
    forall a.
    String
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall e (m :: * -> *) a. MonadError e m => e -> m a
    throwError String
    "Cannot handle invariant allocations in SegBlock."
        SegLevel
    _ ->
          ()
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) ()
    forall a.
    a -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

      if null variant_allocs && null invariant_allocs
        then pure (mempty, (lvl, ops, kbody))
        else do
          (lvl_stms, lvl', grid) <- ensureGridKnown lvl
          allocsForBody variant_allocs invariant_allocs grid space kbody kbody' $
            \RebaseMap
    offsets Stms GPUMem
    alloc_stms KernelBody GPUMem
    kbody'' -> do
              ops'' <- [Lambda GPUMem]
    -> (Lambda GPUMem -> OffsetM (Lambda GPUMem))
    -> OffsetM [Lambda GPUMem]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [Lambda GPUMem]
    ops' ((Lambda GPUMem -> OffsetM (Lambda GPUMem))
     -> OffsetM [Lambda GPUMem])
    -> (Lambda GPUMem -> OffsetM (Lambda GPUMem))
    -> OffsetM [Lambda GPUMem]
    forall a b. (a -> b) -> a -> b
    $ \Lambda GPUMem
    op' ->
                Scope GPUMem -> OffsetM (Lambda GPUMem) -> OffsetM (Lambda GPUMem)
    forall a. Scope GPUMem -> OffsetM a -> OffsetM a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope (Lambda GPUMem -> Scope GPUMem
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf Lambda GPUMem
    op') (OffsetM (Lambda GPUMem) -> OffsetM (Lambda GPUMem))
    -> OffsetM (Lambda GPUMem) -> OffsetM (Lambda GPUMem)
    forall a b. (a -> b) -> a -> b
    $ RebaseMap -> Lambda GPUMem -> OffsetM (Lambda GPUMem)
    offsetMemoryInLambda RebaseMap
    offsets Lambda GPUMem
    op'
              pure (lvl_stms <> alloc_stms, (lvl', ops'', kbody''))
      where
        bound_in_kernel :: Names
    bound_in_kernel =
          [VName] -> Names
    namesFromList (Map VName (NameInfo (ZonkAny 3)) -> [VName]
    forall k a. Map k a -> [k]
    M.keys (Map VName (NameInfo (ZonkAny 3)) -> [VName])
    -> Map VName (NameInfo (ZonkAny 3)) -> [VName]
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> Map VName (NameInfo (ZonkAny 3))
    forall rep. SegSpace -> Scope rep
    scopeOfSegSpace SegSpace
    space)
            Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> KernelBody GPUMem -> Names
    boundInKernelBody KernelBody GPUMem
    kbody

    boundInKernelBody :: KernelBody GPUMem -> Names
    boundInKernelBody :: KernelBody GPUMem -> Names
    boundInKernelBody = [VName] -> Names
    namesFromList ([VName] -> Names)
    -> (KernelBody GPUMem -> [VName]) -> KernelBody GPUMem -> Names
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Scope GPUMem -> [VName]
    forall k a. Map k a -> [k]
    M.keys (Scope GPUMem -> [VName])
    -> (KernelBody GPUMem -> Scope GPUMem)
    -> KernelBody GPUMem
    -> [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Stms GPUMem -> Scope GPUMem
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf (Stms GPUMem -> Scope GPUMem)
    -> (KernelBody GPUMem -> Stms GPUMem)
    -> KernelBody GPUMem
    -> Scope GPUMem
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . KernelBody GPUMem -> Stms GPUMem
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms

    addStmsToKernelBody :: Stms GPUMem -> KernelBody GPUMem -> KernelBody GPUMem
    addStmsToKernelBody :: Stms GPUMem -> KernelBody GPUMem -> KernelBody GPUMem
    addStmsToKernelBody Stms GPUMem
    stms KernelBody GPUMem
    kbody =
      KernelBody GPUMem
    kbody {kernelBodyStms = stms <> kernelBodyStms kbody}

    allocsForBody ::
      Extraction ->
      Extraction ->
      KernelGrid ->
      SegSpace ->
      KernelBody GPUMem ->
      KernelBody GPUMem ->
      (RebaseMap -> Stms GPUMem -> KernelBody GPUMem -> OffsetM b) ->
      ExpandM b
    allocsForBody :: forall b.
    Extraction
    -> Extraction
    -> KernelGrid
    -> SegSpace
    -> KernelBody GPUMem
    -> KernelBody GPUMem
    -> (RebaseMap -> Stms GPUMem -> KernelBody GPUMem -> OffsetM b)
    -> ExpandM b
    allocsForBody Extraction
    variant_allocs Extraction
    invariant_allocs KernelGrid
    grid SegSpace
    space KernelBody GPUMem
    kbody KernelBody GPUMem
    kbody' RebaseMap -> Stms GPUMem -> KernelBody GPUMem -> OffsetM b
    m = do
      (alloc_offsets, alloc_stms) <-
        KernelGrid
    -> SegSpace
    -> Stms GPUMem
    -> Extraction
    -> Extraction
    -> ExpandM (RebaseMap, Stms GPUMem)
    memoryRequirements
          KernelGrid
    grid
          SegSpace
    space
          (KernelBody GPUMem -> Stms GPUMem
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody GPUMem
    kbody)
          Extraction
    variant_allocs
          Extraction
    invariant_allocs

      -- We assume that any shared memory allocations can be inserted back
      -- into kbody'. This would not work if we had SegRed/SegScan
      -- operations that performed shared memory allocations. We don't
      -- currently, and if we would in the future, we would need to be
      -- more careful about summarising the allocations in
      -- transformScanRed.
      let (alloc_stms_dev, alloc_stms_shared) =
            Seq.partition (not . isSharedAlloc) alloc_stms

      scope <- askScope
      let scope' = SegSpace -> Scope GPUMem
    forall rep. SegSpace -> Scope rep
    scopeOfSegSpace SegSpace
    space Scope GPUMem -> Scope GPUMem -> Scope GPUMem
    forall a. Semigroup a => a -> a -> a
    <> Scope GPUMem
    scope Scope GPUMem -> Scope GPUMem -> Scope GPUMem
    forall a. Semigroup a => a -> a -> a
    <> Stms GPUMem -> Scope GPUMem
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf Stms GPUMem
    alloc_stms
      either throwError pure <=< runOffsetM scope' $ do
        kbody'' <-
          addStmsToKernelBody alloc_stms_shared
            <$> offsetMemoryInKernelBody alloc_offsets kbody'
        m alloc_offsets alloc_stms_dev kbody''
      where
        isSharedAlloc :: Stm rep -> Bool
    isSharedAlloc (Let Pat (LetDec rep)
    _ StmAux (ExpDec rep)
    _ (Op (Alloc SubExp
    _ (Space String
    "shared")))) = Bool
    True
        isSharedAlloc Stm rep
    _ = Bool
    False

    memoryRequirements ::
      KernelGrid ->
      SegSpace ->
      Stms GPUMem ->
      Extraction ->
      Extraction ->
      ExpandM (RebaseMap, Stms GPUMem)
    memoryRequirements :: KernelGrid
    -> SegSpace
    -> Stms GPUMem
    -> Extraction
    -> Extraction
    -> ExpandM (RebaseMap, Stms GPUMem)
    memoryRequirements KernelGrid
    grid SegSpace
    space Stms GPUMem
    kstms Extraction
    variant_allocs Extraction
    invariant_allocs = do
      (num_threads, num_threads_stms) <-
        Builder GPUMem SubExp
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         (SubExp, Stms GPUMem)
    forall (m :: * -> *) somerep rep a.
    (MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
    Builder rep a -> m (a, Stms rep)
    runBuilder (Builder GPUMem SubExp
     -> ReaderT
          (Scope GPUMem)
          (StateT VNameSource (Either String))
          (SubExp, Stms GPUMem))
    -> (BasicOp -> Builder GPUMem SubExp)
    -> BasicOp
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         (SubExp, Stms GPUMem)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . String
    -> Exp (Rep (BuilderT GPUMem (State VNameSource)))
    -> Builder GPUMem SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    "num_threads" (Exp GPUMem -> Builder GPUMem SubExp)
    -> (BasicOp -> Exp GPUMem) -> BasicOp -> Builder GPUMem SubExp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . BasicOp -> Exp GPUMem
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp
     -> ReaderT
          (Scope GPUMem)
          (StateT VNameSource (Either String))
          (SubExp, Stms GPUMem))
    -> BasicOp
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         (SubExp, Stms GPUMem)
    forall a b. (a -> b) -> a -> b
    $
          BinOp -> SubExp -> SubExp -> BasicOp
    BinOp
            (IntType -> Overflow -> BinOp
    Mul IntType
    Int64 Overflow
    OverflowUndef)
            (Count NumBlocks SubExp -> SubExp
    forall {k} (u :: k) e. Count u e -> e
    unCount (Count NumBlocks SubExp -> SubExp)
    -> Count NumBlocks SubExp -> SubExp
    forall a b. (a -> b) -> a -> b
    $ KernelGrid -> Count NumBlocks SubExp
    gridNumBlocks KernelGrid
    grid)
            (Count BlockSize SubExp -> SubExp
    forall {k} (u :: k) e. Count u e -> e
    unCount (Count BlockSize SubExp -> SubExp)
    -> Count BlockSize SubExp -> SubExp
    forall a b. (a -> b) -> a -> b
    $ KernelGrid -> Count BlockSize SubExp
    gridBlockSize KernelGrid
    grid)

      (invariant_alloc_stms, invariant_alloc_offsets) <-
        inScopeOf num_threads_stms $
          expandedInvariantAllocations
            num_threads
            (gridNumBlocks grid)
            (gridBlockSize grid)
            invariant_allocs

      (variant_alloc_stms, variant_alloc_offsets) <-
        inScopeOf num_threads_stms $
          expandedVariantAllocations
            num_threads
            space
            kstms
            variant_allocs

      pure
        ( invariant_alloc_offsets <> variant_alloc_offsets,
          num_threads_stms <> invariant_alloc_stms <> variant_alloc_stms
        )

    type Exp64 = TPrimExp Int64 VName

    -- | Identifying the spot where an allocation occurs in terms of its
    -- level and unique thread ID.
    type User = (SegLevel, [Exp64])

    -- | A description of allocations that have been extracted, and how
    -- much memory (and which space) is needed.
    type Extraction = M.Map VName (User, SubExp, Space)

    extractKernelBodyAllocations ::
      User ->
      Names ->
      Names ->
      KernelBody GPUMem ->
      ( KernelBody GPUMem,
        Extraction
      )
    extractKernelBodyAllocations :: (SegLevel, [TPrimExp Int64 VName])
    -> Names
    -> Names
    -> KernelBody GPUMem
    -> (KernelBody GPUMem, Extraction)
    extractKernelBodyAllocations (SegLevel, [TPrimExp Int64 VName])
    lvl Names
    bound_outside Names
    bound_kernel =
      (SegLevel, [TPrimExp Int64 VName])
    -> Names
    -> Names
    -> (KernelBody GPUMem -> Stms GPUMem)
    -> (Stms GPUMem -> KernelBody GPUMem -> KernelBody GPUMem)
    -> KernelBody GPUMem
    -> (KernelBody GPUMem, Extraction)
    forall body.
    (SegLevel, [TPrimExp Int64 VName])
    -> Names
    -> Names
    -> (body -> Stms GPUMem)
    -> (Stms GPUMem -> body -> body)
    -> body
    -> (body, Extraction)
    extractGenericBodyAllocations (SegLevel, [TPrimExp Int64 VName])
    lvl Names
    bound_outside Names
    bound_kernel KernelBody GPUMem -> Stms GPUMem
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms ((Stms GPUMem -> KernelBody GPUMem -> KernelBody GPUMem)
     -> KernelBody GPUMem -> (KernelBody GPUMem, Extraction))
    -> (Stms GPUMem -> KernelBody GPUMem -> KernelBody GPUMem)
    -> KernelBody GPUMem
    -> (KernelBody GPUMem, Extraction)
    forall a b. (a -> b) -> a -> b
    $
        \Stms GPUMem
    stms KernelBody GPUMem
    kbody -> KernelBody GPUMem
    kbody {kernelBodyStms = stms}

    extractBodyAllocations ::
      User ->
      Names ->
      Names ->
      Body GPUMem ->
      (Body GPUMem, Extraction)
    extractBodyAllocations :: (SegLevel, [TPrimExp Int64 VName])
    -> Names -> Names -> Body GPUMem -> (Body GPUMem, Extraction)
    extractBodyAllocations (SegLevel, [TPrimExp Int64 VName])
    user Names
    bound_outside Names
    bound_kernel =
      (SegLevel, [TPrimExp Int64 VName])
    -> Names
    -> Names
    -> (Body GPUMem -> Stms GPUMem)
    -> (Stms GPUMem -> Body GPUMem -> Body GPUMem)
    -> Body GPUMem
    -> (Body GPUMem, Extraction)
    forall body.
    (SegLevel, [TPrimExp Int64 VName])
    -> Names
    -> Names
    -> (body -> Stms GPUMem)
    -> (Stms GPUMem -> body -> body)
    -> body
    -> (body, Extraction)
    extractGenericBodyAllocations (SegLevel, [TPrimExp Int64 VName])
    user Names
    bound_outside Names
    bound_kernel Body GPUMem -> Stms GPUMem
    forall rep. Body rep -> Stms rep
    bodyStms ((Stms GPUMem -> Body GPUMem -> Body GPUMem)
     -> Body GPUMem -> (Body GPUMem, Extraction))
    -> (Stms GPUMem -> Body GPUMem -> Body GPUMem)
    -> Body GPUMem
    -> (Body GPUMem, Extraction)
    forall a b. (a -> b) -> a -> b
    $
        \Stms GPUMem
    stms Body GPUMem
    body -> Body GPUMem
    body {bodyStms = stms}

    extractLambdaAllocations ::
      User ->
      Names ->
      Names ->
      Lambda GPUMem ->
      (Lambda GPUMem, Extraction)
    extractLambdaAllocations :: (SegLevel, [TPrimExp Int64 VName])
    -> Names -> Names -> Lambda GPUMem -> (Lambda GPUMem, Extraction)
    extractLambdaAllocations (SegLevel, [TPrimExp Int64 VName])
    user Names
    bound_outside Names
    bound_kernel Lambda GPUMem
    lam =
      (Lambda GPUMem
    lam {lambdaBody = body'}, Extraction
    allocs)
      where
        (Body GPUMem
    body', Extraction
    allocs) =
          (SegLevel, [TPrimExp Int64 VName])
    -> Names -> Names -> Body GPUMem -> (Body GPUMem, Extraction)
    extractBodyAllocations (SegLevel, [TPrimExp Int64 VName])
    user Names
    bound_outside Names
    bound_kernel (Body GPUMem -> (Body GPUMem, Extraction))
    -> Body GPUMem -> (Body GPUMem, Extraction)
    forall a b. (a -> b) -> a -> b
    $
            Lambda GPUMem -> Body GPUMem
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda GPUMem
    lam

    extractGenericBodyAllocations ::
      User ->
      Names ->
      Names ->
      (body -> Stms GPUMem) ->
      (Stms GPUMem -> body -> body) ->
      body ->
      ( body,
        Extraction
      )
    extractGenericBodyAllocations :: forall body.
    (SegLevel, [TPrimExp Int64 VName])
    -> Names
    -> Names
    -> (body -> Stms GPUMem)
    -> (Stms GPUMem -> body -> body)
    -> body
    -> (body, Extraction)
    extractGenericBodyAllocations (SegLevel, [TPrimExp Int64 VName])
    user Names
    bound_outside Names
    bound_kernel body -> Stms GPUMem
    get_stms Stms GPUMem -> body -> body
    set_stms body
    body =
      let bound_kernel' :: Names
    bound_kernel' = Names
    bound_kernel Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Stms GPUMem -> Names
    forall rep. Stms rep -> Names
    boundByStms (body -> Stms GPUMem
    get_stms body
    body)
          ([Stm GPUMem]
    stms, Extraction
    allocs) =
            Writer Extraction [Stm GPUMem] -> ([Stm GPUMem], Extraction)
    forall w a. Writer w a -> (a, w)
    runWriter (Writer Extraction [Stm GPUMem] -> ([Stm GPUMem], Extraction))
    -> (WriterT Extraction Identity [Maybe (Stm GPUMem)]
        -> Writer Extraction [Stm GPUMem])
    -> WriterT Extraction Identity [Maybe (Stm GPUMem)]
    -> ([Stm GPUMem], Extraction)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ([Maybe (Stm GPUMem)] -> [Stm GPUMem])
    -> WriterT Extraction Identity [Maybe (Stm GPUMem)]
    -> Writer Extraction [Stm GPUMem]
    forall a b.
    (a -> b)
    -> WriterT Extraction Identity a -> WriterT Extraction Identity b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [Maybe (Stm GPUMem)] -> [Stm GPUMem]
    forall a. [Maybe a] -> [a]
    catMaybes (WriterT Extraction Identity [Maybe (Stm GPUMem)]
     -> ([Stm GPUMem], Extraction))
    -> WriterT Extraction Identity [Maybe (Stm GPUMem)]
    -> ([Stm GPUMem], Extraction)
    forall a b. (a -> b) -> a -> b
    $
              (Stm GPUMem -> WriterT Extraction Identity (Maybe (Stm GPUMem)))
    -> [Stm GPUMem] -> WriterT Extraction Identity [Maybe (Stm GPUMem)]
    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 ((SegLevel, [TPrimExp Int64 VName])
    -> Names
    -> Names
    -> Stm GPUMem
    -> WriterT Extraction Identity (Maybe (Stm GPUMem))
    extractStmAllocations (SegLevel, [TPrimExp Int64 VName])
    user Names
    bound_outside Names
    bound_kernel') ([Stm GPUMem] -> WriterT Extraction Identity [Maybe (Stm GPUMem)])
    -> [Stm GPUMem] -> WriterT Extraction Identity [Maybe (Stm GPUMem)]
    forall a b. (a -> b) -> a -> b
    $
                Stms GPUMem -> [Stm GPUMem]
    forall rep. Stms rep -> [Stm rep]
    stmsToList (body -> Stms GPUMem
    get_stms body
    body)
       in (Stms GPUMem -> body -> body
    set_stms ([Stm GPUMem] -> Stms GPUMem
    forall rep. [Stm rep] -> Stms rep
    stmsFromList [Stm GPUMem]
    stms) body
    body, Extraction
    allocs)

    expandable :: User -> Space -> Bool
    expandable :: (SegLevel, [TPrimExp Int64 VName]) -> Space -> Bool
    expandable (SegBlock {}, [TPrimExp Int64 VName]
    _) (Space String
    "shared") = Bool
    False
    expandable (SegLevel, [TPrimExp Int64 VName])
    _ ScalarSpace {} = Bool
    False
    expandable (SegLevel, [TPrimExp Int64 VName])
    _ Space
    _ = Bool
    True

    notScalar :: Space -> Bool
    notScalar :: Space -> Bool
    notScalar ScalarSpace {} = Bool
    False
    notScalar Space
    _ = Bool
    True

    extractStmAllocations ::
      User ->
      Names ->
      Names ->
      Stm GPUMem ->
      Writer Extraction (Maybe (Stm GPUMem))
    extractStmAllocations :: (SegLevel, [TPrimExp Int64 VName])
    -> Names
    -> Names
    -> Stm GPUMem
    -> WriterT Extraction Identity (Maybe (Stm GPUMem))
    extractStmAllocations (SegLevel, [TPrimExp Int64 VName])
    user Names
    bound_outside Names
    bound_kernel (Let (Pat [PatElem (LetDec GPUMem)
    patElem]) StmAux (ExpDec GPUMem)
    _ (Op (Alloc SubExp
    size Space
    space)))
      | (SegLevel, [TPrimExp Int64 VName]) -> Space -> Bool
    expandable (SegLevel, [TPrimExp Int64 VName])
    user Space
    space Bool -> Bool -> Bool
    && SubExp -> Bool
    expandableSize SubExp
    size
          -- FIXME: the '&& notScalar space' part is a hack because we
          -- don't otherwise hoist the sizes out far enough, and we
          -- promise to be super-duper-careful about not having variant
          -- scalar allocations.
          Bool -> Bool -> Bool
    || (SubExp -> Bool
    boundInKernel SubExp
    size Bool -> Bool -> Bool
    && Space -> Bool
    notScalar Space
    space) = do
          Extraction -> WriterT Extraction Identity ()
    forall w (m :: * -> *). MonadWriter w m => w -> m ()
    tell (Extraction -> WriterT Extraction Identity ())
    -> Extraction -> WriterT Extraction Identity ()
    forall a b. (a -> b) -> a -> b
    $ VName
    -> ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space)
    -> Extraction
    forall k a. k -> a -> Map k a
    M.singleton (PatElem (MemInfo SubExp NoUniqueness MemBind) -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem (LetDec GPUMem)
    PatElem (MemInfo SubExp NoUniqueness MemBind)
    patElem) ((SegLevel, [TPrimExp Int64 VName])
    user, SubExp
    size, Space
    space)
          Maybe (Stm GPUMem)
    -> WriterT Extraction Identity (Maybe (Stm GPUMem))
    forall a. a -> WriterT Extraction Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe (Stm GPUMem)
    forall a. Maybe a
    Nothing
      where
        expandableSize :: SubExp -> Bool
    expandableSize (Var VName
    v) = VName
    v VName -> Names -> Bool
    `nameIn` Names
    bound_outside Bool -> Bool -> Bool
    || VName
    v VName -> Names -> Bool
    `nameIn` Names
    bound_kernel
        expandableSize Constant {} = Bool
    True
        boundInKernel :: SubExp -> Bool
    boundInKernel (Var VName
    v) = VName
    v VName -> Names -> Bool
    `nameIn` Names
    bound_kernel
        boundInKernel Constant {} = Bool
    False
    extractStmAllocations (SegLevel, [TPrimExp Int64 VName])
    user Names
    bound_outside Names
    bound_kernel Stm GPUMem
    stm = do
      e <- Mapper GPUMem GPUMem (WriterT Extraction Identity)
    -> Exp GPUMem -> WriterT Extraction Identity (Exp GPUMem)
    forall (m :: * -> *) frep trep.
    Monad m =>
    Mapper frep trep m -> Exp frep -> m (Exp trep)
    mapExpM ((SegLevel, [TPrimExp Int64 VName])
    -> Mapper GPUMem GPUMem (WriterT Extraction Identity)
    expMapper (SegLevel, [TPrimExp Int64 VName])
    user) (Exp GPUMem -> WriterT Extraction Identity (Exp GPUMem))
    -> Exp GPUMem -> WriterT Extraction Identity (Exp GPUMem)
    forall a b. (a -> b) -> a -> b
    $ Stm GPUMem -> Exp GPUMem
    forall rep. Stm rep -> Exp rep
    stmExp Stm GPUMem
    stm
      pure $ Just $ stm {stmExp = e}
      where
        expMapper :: (SegLevel, [TPrimExp Int64 VName])
    -> Mapper GPUMem GPUMem (WriterT Extraction Identity)
    expMapper (SegLevel, [TPrimExp Int64 VName])
    user' =
          (forall rep (m :: * -> *). Monad m => Mapper rep rep m
    identityMapper @GPUMem)
            { mapOnBody = const $ onBody user',
              mapOnOp = onOp user'
            }

        onBody :: (SegLevel, [TPrimExp Int64 VName])
    -> Body GPUMem -> WriterT Extraction Identity (Body GPUMem)
    onBody (SegLevel, [TPrimExp Int64 VName])
    user' Body GPUMem
    body = do
          let (Body GPUMem
    body', Extraction
    allocs) = (SegLevel, [TPrimExp Int64 VName])
    -> Names -> Names -> Body GPUMem -> (Body GPUMem, Extraction)
    extractBodyAllocations (SegLevel, [TPrimExp Int64 VName])
    user' Names
    bound_outside Names
    bound_kernel Body GPUMem
    body
          Extraction -> WriterT Extraction Identity ()
    forall w (m :: * -> *). MonadWriter w m => w -> m ()
    tell Extraction
    allocs
          Body GPUMem -> WriterT Extraction Identity (Body GPUMem)
    forall a. a -> WriterT Extraction Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Body GPUMem
    body'

        onOp :: (SegLevel, [TPrimExp Int64 VName])
    -> MemOp (HostOp NoOp) GPUMem
    -> WriterT Extraction Identity (MemOp (HostOp NoOp) GPUMem)
    onOp (SegLevel
    _, [TPrimExp Int64 VName]
    user_ids) (Inner (SegOp SegOp SegLevel GPUMem
    op)) =
          HostOp NoOp GPUMem -> MemOp (HostOp NoOp) GPUMem
    forall (inner :: * -> *) rep. inner rep -> MemOp inner rep
    Inner (HostOp NoOp GPUMem -> MemOp (HostOp NoOp) GPUMem)
    -> (SegOp SegLevel GPUMem -> HostOp NoOp GPUMem)
    -> SegOp SegLevel GPUMem
    -> MemOp (HostOp NoOp) GPUMem
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegOp SegLevel GPUMem -> HostOp NoOp GPUMem
    forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
    SegOp (SegOp SegLevel GPUMem -> MemOp (HostOp NoOp) GPUMem)
    -> WriterT Extraction Identity (SegOp SegLevel GPUMem)
    -> WriterT Extraction Identity (MemOp (HostOp NoOp) GPUMem)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SegOpMapper SegLevel GPUMem GPUMem (WriterT Extraction Identity)
    -> SegOp SegLevel GPUMem
    -> WriterT Extraction Identity (SegOp SegLevel GPUMem)
    forall (m :: * -> *) lvl frep trep.
    Monad m =>
    SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
    mapSegOpM ((SegLevel, [TPrimExp Int64 VName])
    -> SegOpMapper SegLevel GPUMem GPUMem (WriterT Extraction Identity)
    opMapper (SegLevel, [TPrimExp Int64 VName])
    user'') SegOp SegLevel GPUMem
    op
          where
            user'' :: (SegLevel, [TPrimExp Int64 VName])
    user'' =
              (SegOp SegLevel GPUMem -> SegLevel
    forall lvl rep. SegOp lvl rep -> lvl
    segLevel SegOp SegLevel GPUMem
    op, [TPrimExp Int64 VName]
    user_ids [TPrimExp Int64 VName]
    -> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
    forall a. [a] -> [a] -> [a]
    ++ [VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 (SegSpace -> VName
    segFlat (SegOp SegLevel GPUMem -> SegSpace
    forall lvl rep. SegOp lvl rep -> SegSpace
    segSpace SegOp SegLevel GPUMem
    op))])
        onOp (SegLevel, [TPrimExp Int64 VName])
    _ MemOp (HostOp NoOp) GPUMem
    op = MemOp (HostOp NoOp) GPUMem
    -> WriterT Extraction Identity (MemOp (HostOp NoOp) GPUMem)
    forall a. a -> WriterT Extraction Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure MemOp (HostOp NoOp) GPUMem
    op

        opMapper :: (SegLevel, [TPrimExp Int64 VName])
    -> SegOpMapper SegLevel GPUMem GPUMem (WriterT Extraction Identity)
    opMapper (SegLevel, [TPrimExp Int64 VName])
    user' =
          SegOpMapper
      SegLevel (ZonkAny 2) (ZonkAny 2) (WriterT Extraction Identity)
    forall (m :: * -> *) lvl rep. Monad m => SegOpMapper lvl rep rep m
    identitySegOpMapper
            { mapOnSegOpLambda = onLambda user',
              mapOnSegOpBody = onKernelBody user'
            }

        onKernelBody :: (SegLevel, [TPrimExp Int64 VName])
    -> KernelBody GPUMem
    -> WriterT Extraction Identity (KernelBody GPUMem)
    onKernelBody (SegLevel, [TPrimExp Int64 VName])
    user' KernelBody GPUMem
    body = do
          let (KernelBody GPUMem
    body', Extraction
    allocs) =
                (SegLevel, [TPrimExp Int64 VName])
    -> Names
    -> Names
    -> KernelBody GPUMem
    -> (KernelBody GPUMem, Extraction)
    extractKernelBodyAllocations (SegLevel, [TPrimExp Int64 VName])
    user' Names
    bound_outside Names
    bound_kernel KernelBody GPUMem
    body
          Extraction -> WriterT Extraction Identity ()
    forall w (m :: * -> *). MonadWriter w m => w -> m ()
    tell Extraction
    allocs
          KernelBody GPUMem
    -> WriterT Extraction Identity (KernelBody GPUMem)
    forall a. a -> WriterT Extraction Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure KernelBody GPUMem
    body'

        onLambda :: (SegLevel, [TPrimExp Int64 VName])
    -> Lambda GPUMem -> WriterT Extraction Identity (Lambda GPUMem)
    onLambda (SegLevel, [TPrimExp Int64 VName])
    user' Lambda GPUMem
    lam = do
          body <- (SegLevel, [TPrimExp Int64 VName])
    -> Body GPUMem -> WriterT Extraction Identity (Body GPUMem)
    onBody (SegLevel, [TPrimExp Int64 VName])
    user' (Body GPUMem -> WriterT Extraction Identity (Body GPUMem))
    -> Body GPUMem -> WriterT Extraction Identity (Body GPUMem)
    forall a b. (a -> b) -> a -> b
    $ Lambda GPUMem -> Body GPUMem
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda GPUMem
    lam
          pure lam {lambdaBody = body}

    genericExpandedInvariantAllocations ::
      (User -> Space -> (Shape, [Exp64])) -> Extraction -> ExpandM (Stms GPUMem, RebaseMap)
    genericExpandedInvariantAllocations :: ((SegLevel, [TPrimExp Int64 VName])
     -> Space -> (ShapeBase SubExp, [TPrimExp Int64 VName]))
    -> Extraction -> ExpandM (Stms GPUMem, RebaseMap)
    genericExpandedInvariantAllocations (SegLevel, [TPrimExp Int64 VName])
    -> Space -> (ShapeBase SubExp, [TPrimExp Int64 VName])
    getNumUsers Extraction
    invariant_allocs = do
      -- We expand the invariant allocations by adding an inner dimension
      -- equal to the number of kernel threads.
      (rebases, alloc_stms) <- Builder GPUMem [RebaseMap]
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         ([RebaseMap], Stms GPUMem)
    forall (m :: * -> *) somerep rep a.
    (MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
    Builder rep a -> m (a, Stms rep)
    runBuilder (Builder GPUMem [RebaseMap]
     -> ReaderT
          (Scope GPUMem)
          (StateT VNameSource (Either String))
          ([RebaseMap], Stms GPUMem))
    -> Builder GPUMem [RebaseMap]
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         ([RebaseMap], Stms GPUMem)
    forall a b. (a -> b) -> a -> b
    $ ((VName, ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space))
     -> BuilderT GPUMem (State VNameSource) RebaseMap)
    -> [(VName, ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space))]
    -> Builder GPUMem [RebaseMap]
    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, ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space))
    -> BuilderT GPUMem (State VNameSource) RebaseMap
    expand ([(VName, ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space))]
     -> Builder GPUMem [RebaseMap])
    -> [(VName, ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space))]
    -> Builder GPUMem [RebaseMap]
    forall a b. (a -> b) -> a -> b
    $ Extraction
    -> [(VName, ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space))]
    forall k a. Map k a -> [(k, a)]
    M.toList Extraction
    invariant_allocs

      pure (alloc_stms, mconcat rebases)
      where
        expand :: (VName, ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space))
    -> BuilderT GPUMem (State VNameSource) RebaseMap
    expand (VName
    mem, ((SegLevel, [TPrimExp Int64 VName])
    user, SubExp
    per_thread_size, Space
    space)) = do
          let num_users :: ShapeBase SubExp
    num_users = (ShapeBase SubExp, [TPrimExp Int64 VName]) -> ShapeBase SubExp
    forall a b. (a, b) -> a
    fst ((ShapeBase SubExp, [TPrimExp Int64 VName]) -> ShapeBase SubExp)
    -> (ShapeBase SubExp, [TPrimExp Int64 VName]) -> ShapeBase SubExp
    forall a b. (a -> b) -> a -> b
    $ (SegLevel, [TPrimExp Int64 VName])
    -> Space -> (ShapeBase SubExp, [TPrimExp Int64 VName])
    getNumUsers (SegLevel, [TPrimExp Int64 VName])
    user Space
    space
              allocpat :: Pat (MemInfo SubExp NoUniqueness MemBind)
    allocpat = [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    -> Pat (MemInfo SubExp NoUniqueness MemBind)
    forall dec. [PatElem dec] -> Pat dec
    Pat [VName
    -> MemInfo SubExp NoUniqueness MemBind
    -> PatElem (MemInfo SubExp NoUniqueness MemBind)
    forall dec. VName -> dec -> PatElem dec
    PatElem VName
    mem (MemInfo SubExp NoUniqueness MemBind
     -> PatElem (MemInfo SubExp NoUniqueness MemBind))
    -> MemInfo SubExp NoUniqueness MemBind
    -> PatElem (MemInfo SubExp NoUniqueness MemBind)
    forall a b. (a -> b) -> a -> b
    $ Space -> MemInfo SubExp NoUniqueness MemBind
    forall d u ret. Space -> MemInfo d u ret
    MemMem Space
    space]
          total_size <-
            String
    -> Exp (Rep (BuilderT GPUMem (State VNameSource)))
    -> BuilderT GPUMem (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m VName
    letExp String
    "total_size" (Exp GPUMem -> BuilderT GPUMem (State VNameSource) VName)
    -> ([TPrimExp Int64 VName]
        -> BuilderT GPUMem (State VNameSource) (Exp GPUMem))
    -> [TPrimExp Int64 VName]
    -> BuilderT GPUMem (State VNameSource) VName
    forall (m :: * -> *) b c a.
    Monad m =>
    (b -> m c) -> (a -> m b) -> a -> m c
    <=< TPrimExp Int64 VName
    -> BuilderT
         GPUMem
         (State VNameSource)
         (Exp (Rep (BuilderT GPUMem (State VNameSource))))
    TPrimExp Int64 VName
    -> BuilderT GPUMem (State VNameSource) (Exp GPUMem)
    forall a (m :: * -> *).
    (ToExp a, MonadBuilder m) =>
    a -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    TPrimExp Int64 VName -> m (Exp (Rep m))
    toExp (TPrimExp Int64 VName
     -> BuilderT GPUMem (State VNameSource) (Exp GPUMem))
    -> ([TPrimExp Int64 VName] -> TPrimExp Int64 VName)
    -> [TPrimExp Int64 VName]
    -> BuilderT GPUMem (State VNameSource) (Exp GPUMem)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [TPrimExp Int64 VName] -> TPrimExp Int64 VName
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    product ([TPrimExp Int64 VName]
     -> BuilderT GPUMem (State VNameSource) VName)
    -> [TPrimExp Int64 VName]
    -> BuilderT GPUMem (State VNameSource) VName
    forall a b. (a -> b) -> a -> b
    $
              SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    per_thread_size TPrimExp Int64 VName
    -> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
    forall a. a -> [a] -> [a]
    : (SubExp -> TPrimExp Int64 VName)
    -> [SubExp] -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> TPrimExp Int64 VName
    pe64 (ShapeBase SubExp -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims ShapeBase SubExp
    num_users)
          letBind allocpat $ Op $ Alloc (Var total_size) space
          pure $ M.singleton mem $ newBase user space

        newBaseThread :: (SegLevel, [TPrimExp Int64 VName])
    -> Space
    -> [TPrimExp Int64 VName]
    -> (TPrimExp Int64 VName, TPrimExp Int64 VName)
    newBaseThread (SegLevel, [TPrimExp Int64 VName])
    user Space
    space [TPrimExp Int64 VName]
    _old_shape =
          let (ShapeBase SubExp
    users_shape, [TPrimExp Int64 VName]
    user_ids) = (SegLevel, [TPrimExp Int64 VName])
    -> Space -> (ShapeBase SubExp, [TPrimExp Int64 VName])
    getNumUsers (SegLevel, [TPrimExp Int64 VName])
    user Space
    space
              dims :: [TPrimExp Int64 VName]
    dims = (SubExp -> TPrimExp Int64 VName)
    -> [SubExp] -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> TPrimExp Int64 VName
    pe64 (ShapeBase SubExp -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims ShapeBase SubExp
    users_shape)
           in ( [TPrimExp Int64 VName]
    -> [TPrimExp Int64 VName] -> TPrimExp Int64 VName
    forall num. IntegralExp num => [num] -> [num] -> num
    flattenIndex [TPrimExp Int64 VName]
    dims [TPrimExp Int64 VName]
    user_ids,
                [TPrimExp Int64 VName] -> TPrimExp Int64 VName
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    product [TPrimExp Int64 VName]
    dims
              )

        newBase :: (SegLevel, [TPrimExp Int64 VName])
    -> Space
    -> [TPrimExp Int64 VName]
    -> (TPrimExp Int64 VName, TPrimExp Int64 VName)
    newBase user :: (SegLevel, [TPrimExp Int64 VName])
    user@(SegThreadInBlock {}, [TPrimExp Int64 VName]
    _) Space
    space = (SegLevel, [TPrimExp Int64 VName])
    -> Space
    -> [TPrimExp Int64 VName]
    -> (TPrimExp Int64 VName, TPrimExp Int64 VName)
    newBaseThread (SegLevel, [TPrimExp Int64 VName])
    user Space
    space
        newBase user :: (SegLevel, [TPrimExp Int64 VName])
    user@(SegThread {}, [TPrimExp Int64 VName]
    _) Space
    space = (SegLevel, [TPrimExp Int64 VName])
    -> Space
    -> [TPrimExp Int64 VName]
    -> (TPrimExp Int64 VName, TPrimExp Int64 VName)
    newBaseThread (SegLevel, [TPrimExp Int64 VName])
    user Space
    space
        newBase user :: (SegLevel, [TPrimExp Int64 VName])
    user@(SegBlock {}, [TPrimExp Int64 VName]
    _) Space
    space = \[TPrimExp Int64 VName]
    _old_shape ->
          let (ShapeBase SubExp
    users_shape, [TPrimExp Int64 VName]
    user_ids) = (SegLevel, [TPrimExp Int64 VName])
    -> Space -> (ShapeBase SubExp, [TPrimExp Int64 VName])
    getNumUsers (SegLevel, [TPrimExp Int64 VName])
    user Space
    space
              dims :: [TPrimExp Int64 VName]
    dims = (SubExp -> TPrimExp Int64 VName)
    -> [SubExp] -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> TPrimExp Int64 VName
    pe64 (ShapeBase SubExp -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims ShapeBase SubExp
    users_shape)
           in ( [TPrimExp Int64 VName]
    -> [TPrimExp Int64 VName] -> TPrimExp Int64 VName
    forall num. IntegralExp num => [num] -> [num] -> num
    flattenIndex [TPrimExp Int64 VName]
    dims [TPrimExp Int64 VName]
    user_ids,
                [TPrimExp Int64 VName] -> TPrimExp Int64 VName
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    product [TPrimExp Int64 VName]
    dims
              )

    expandedInvariantAllocations ::
      SubExp ->
      Count NumBlocks SubExp ->
      Count BlockSize SubExp ->
      Extraction ->
      ExpandM (Stms GPUMem, RebaseMap)
    expandedInvariantAllocations :: SubExp
    -> Count NumBlocks SubExp
    -> Count BlockSize SubExp
    -> Extraction
    -> ExpandM (Stms GPUMem, RebaseMap)
    expandedInvariantAllocations SubExp
    num_threads (Count SubExp
    num_tblocks) (Count SubExp
    tblock_size) =
      ((SegLevel, [TPrimExp Int64 VName])
     -> Space -> (ShapeBase SubExp, [TPrimExp Int64 VName]))
    -> Extraction -> ExpandM (Stms GPUMem, RebaseMap)
    genericExpandedInvariantAllocations (SegLevel, [TPrimExp Int64 VName])
    -> Space -> (ShapeBase SubExp, [TPrimExp Int64 VName])
    getNumUsers
      where
        getNumUsers :: (SegLevel, [TPrimExp Int64 VName])
    -> Space -> (ShapeBase SubExp, [TPrimExp Int64 VName])
    getNumUsers (SegThread {}, [TPrimExp Int64 VName
    gtid]) Space
    _ = ([SubExp] -> ShapeBase SubExp
    forall d. [d] -> ShapeBase d
    Shape [SubExp
    num_threads], [TPrimExp Int64 VName
    gtid])
        getNumUsers (SegThread {}, [TPrimExp Int64 VName
    gid, TPrimExp Int64 VName
    ltid]) Space
    _ = ([SubExp] -> ShapeBase SubExp
    forall d. [d] -> ShapeBase d
    Shape [SubExp
    num_tblocks, SubExp
    tblock_size], [TPrimExp Int64 VName
    gid, TPrimExp Int64 VName
    ltid])
        getNumUsers (SegThreadInBlock {}, [TPrimExp Int64 VName
    gtid]) Space
    _ = ([SubExp] -> ShapeBase SubExp
    forall d. [d] -> ShapeBase d
    Shape [SubExp
    num_threads], [TPrimExp Int64 VName
    gtid])
        getNumUsers (SegThreadInBlock {}, [TPrimExp Int64 VName
    _gid, TPrimExp Int64 VName
    ltid]) (Space String
    "shared") =
          ([SubExp] -> ShapeBase SubExp
    forall d. [d] -> ShapeBase d
    Shape [SubExp
    tblock_size], [TPrimExp Int64 VName
    ltid])
        getNumUsers (SegThreadInBlock {}, [TPrimExp Int64 VName
    gid, TPrimExp Int64 VName
    ltid]) (Space String
    "device") =
          ([SubExp] -> ShapeBase SubExp
    forall d. [d] -> ShapeBase d
    Shape [SubExp
    num_tblocks, SubExp
    tblock_size], [TPrimExp Int64 VName
    gid, TPrimExp Int64 VName
    ltid])
        getNumUsers (SegBlock {}, [TPrimExp Int64 VName
    gid]) Space
    _ = ([SubExp] -> ShapeBase SubExp
    forall d. [d] -> ShapeBase d
    Shape [SubExp
    num_tblocks], [TPrimExp Int64 VName
    gid])
        getNumUsers (SegLevel, [TPrimExp Int64 VName])
    user Space
    space = String -> (ShapeBase SubExp, [TPrimExp Int64 VName])
    forall a. HasCallStack => String -> a
    error (String -> (ShapeBase SubExp, [TPrimExp Int64 VName]))
    -> String -> (ShapeBase SubExp, [TPrimExp Int64 VName])
    forall a b. (a -> b) -> a -> b
    $ String
    "getNumUsers: unhandled " String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ ((SegLevel, [TPrimExp Int64 VName]), Space) -> String
    forall a. Show a => a -> String
    show ((SegLevel, [TPrimExp Int64 VName])
    user, Space
    space)

    expandedVariantAllocations ::
      SubExp ->
      SegSpace ->
      Stms GPUMem ->
      Extraction ->
      ExpandM (Stms GPUMem, RebaseMap)
    expandedVariantAllocations :: SubExp
    -> SegSpace
    -> Stms GPUMem
    -> Extraction
    -> ExpandM (Stms GPUMem, RebaseMap)
    expandedVariantAllocations SubExp
    _ SegSpace
    _ Stms GPUMem
    _ Extraction
    variant_allocs
      | Extraction -> Bool
    forall a. Map VName a -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null Extraction
    variant_allocs = (Stms GPUMem, RebaseMap) -> ExpandM (Stms GPUMem, RebaseMap)
    forall a.
    a -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Stms GPUMem
    forall a. Monoid a => a
    mempty, RebaseMap
    forall a. Monoid a => a
    mempty)
    expandedVariantAllocations SubExp
    num_threads SegSpace
    kspace Stms GPUMem
    kstms Extraction
    variant_allocs = do
      let sizes_to_blocks :: [(SubExp, [(VName, Space)])]
    sizes_to_blocks = Extraction -> [(SubExp, [(VName, Space)])]
    removeCommonSizes Extraction
    variant_allocs
          variant_sizes :: [SubExp]
    variant_sizes = ((SubExp, [(VName, Space)]) -> SubExp)
    -> [(SubExp, [(VName, Space)])] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map (SubExp, [(VName, Space)]) -> SubExp
    forall a b. (a, b) -> a
    fst [(SubExp, [(VName, Space)])]
    sizes_to_blocks

      (slice_stms, offsets, size_sums) <-
        SubExp
    -> [SubExp]
    -> SegSpace
    -> Stms GPUMem
    -> ExpandM (Stms GPU, [VName], [VName])
    sliceKernelSizes SubExp
    num_threads [SubExp]
    variant_sizes SegSpace
    kspace Stms GPUMem
    kstms
      -- Note the recursive call to expand allocations inside the newly
      -- produced kernels.
      slice_stms_tmp <- simplifyStms =<< explicitAllocationsInStms slice_stms
      slice_stms' <- transformStms slice_stms_tmp

      let variant_allocs' :: [(VName, (SubExp, SubExp, Space))]
          variant_allocs' =
            [[(VName, (SubExp, SubExp, Space))]]
    -> [(VName, (SubExp, SubExp, Space))]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat ([[(VName, (SubExp, SubExp, Space))]]
     -> [(VName, (SubExp, SubExp, Space))])
    -> [[(VName, (SubExp, SubExp, Space))]]
    -> [(VName, (SubExp, SubExp, Space))]
    forall a b. (a -> b) -> a -> b
    $
              ([(VName, Space)]
     -> (VName, VName) -> [(VName, (SubExp, SubExp, Space))])
    -> [[(VName, Space)]]
    -> [(VName, VName)]
    -> [[(VName, (SubExp, SubExp, Space))]]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith [(VName, Space)]
    -> (VName, VName) -> [(VName, (SubExp, SubExp, Space))]
    forall {a} {c}.
    [(a, c)] -> (VName, VName) -> [(a, (SubExp, SubExp, c))]
    memInfo (((SubExp, [(VName, Space)]) -> [(VName, Space)])
    -> [(SubExp, [(VName, Space)])] -> [[(VName, Space)]]
    forall a b. (a -> b) -> [a] -> [b]
    map (SubExp, [(VName, Space)]) -> [(VName, Space)]
    forall a b. (a, b) -> b
    snd [(SubExp, [(VName, Space)])]
    sizes_to_blocks) ([VName] -> [VName] -> [(VName, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    offsets [VName]
    size_sums)
          memInfo [(a, c)]
    blocks (VName
    offset, VName
    total_size) =
            [(a
    mem, (VName -> SubExp
    Var VName
    offset, VName -> SubExp
    Var VName
    total_size, c
    space)) | (a
    mem, c
    space) <- [(a, c)]
    blocks]

      -- We expand the invariant allocations by adding an inner dimension
      -- equal to the sum of the sizes required by different threads.
      (alloc_stms, rebases) <- mapAndUnzipM expand variant_allocs'

      pure (slice_stms' <> stmsFromList alloc_stms, mconcat rebases)
      where
        expand :: (VName, (SubExp, SubExp, Space))
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         (Stm GPUMem, RebaseMap)
    expand (VName
    mem, (SubExp
    _offset, SubExp
    total_size, Space
    space)) = do
          let allocpat :: Pat (MemInfo SubExp NoUniqueness MemBind)
    allocpat = [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    -> Pat (MemInfo SubExp NoUniqueness MemBind)
    forall dec. [PatElem dec] -> Pat dec
    Pat [VName
    -> MemInfo SubExp NoUniqueness MemBind
    -> PatElem (MemInfo SubExp NoUniqueness MemBind)
    forall dec. VName -> dec -> PatElem dec
    PatElem VName
    mem (MemInfo SubExp NoUniqueness MemBind
     -> PatElem (MemInfo SubExp NoUniqueness MemBind))
    -> MemInfo SubExp NoUniqueness MemBind
    -> PatElem (MemInfo SubExp NoUniqueness MemBind)
    forall a b. (a -> b) -> a -> b
    $ Space -> MemInfo SubExp NoUniqueness MemBind
    forall d u ret. Space -> MemInfo d u ret
    MemMem Space
    space]
          (Stm GPUMem, RebaseMap)
    -> ReaderT
         (Scope GPUMem)
         (StateT VNameSource (Either String))
         (Stm GPUMem, RebaseMap)
    forall a.
    a -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
            ( Pat (LetDec GPUMem)
    -> StmAux (ExpDec GPUMem) -> Exp GPUMem -> Stm GPUMem
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let Pat (LetDec GPUMem)
    Pat (MemInfo SubExp NoUniqueness MemBind)
    allocpat (() -> StmAux ()
    forall dec. dec -> StmAux dec
    defAux ()) (Exp GPUMem -> Stm GPUMem) -> Exp GPUMem -> Stm GPUMem
    forall a b. (a -> b) -> a -> b
    $ Op GPUMem -> Exp GPUMem
    forall rep. Op rep -> Exp rep
    Op (Op GPUMem -> Exp GPUMem) -> Op GPUMem -> Exp GPUMem
    forall a b. (a -> b) -> a -> b
    $ SubExp -> Space -> MemOp (HostOp NoOp) GPUMem
    forall (inner :: * -> *) rep. SubExp -> Space -> MemOp inner rep
    Alloc SubExp
    total_size Space
    space,
              VName
    -> ([TPrimExp Int64 VName]
        -> (TPrimExp Int64 VName, TPrimExp Int64 VName))
    -> RebaseMap
    forall k a. k -> a -> Map k a
    M.singleton VName
    mem [TPrimExp Int64 VName]
    -> (TPrimExp Int64 VName, TPrimExp Int64 VName)
    newBase
            )

        num_threads' :: TPrimExp Int64 VName
    num_threads' = SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    num_threads
        gtid :: TPrimExp Int64 VName
    gtid = VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 (VName -> TPrimExp Int64 VName) -> VName -> TPrimExp Int64 VName
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> VName
    segFlat SegSpace
    kspace

        -- For the variant allocations, we add an inner dimension,
        -- which is then offset by a thread-specific amount.
        newBase :: [TPrimExp Int64 VName]
    -> (TPrimExp Int64 VName, TPrimExp Int64 VName)
    newBase [TPrimExp Int64 VName]
    _old_shape =
          (TPrimExp Int64 VName
    gtid, TPrimExp Int64 VName
    num_threads')

    type Expansion = (Exp64, Exp64)

    -- | A map from memory block names to index function embeddings..
    type RebaseMap = M.Map VName ([Exp64] -> Expansion)

    --- Modifying the index functions of code.

    newtype OffsetM a
      = OffsetM (BuilderT GPUMem (StateT VNameSource (Either String)) a)
      deriving
        ( Functor OffsetM
    Functor OffsetM =>
    (forall a. a -> OffsetM a)
    -> (forall a b. OffsetM (a -> b) -> OffsetM a -> OffsetM b)
    -> (forall a b c.
        (a -> b -> c) -> OffsetM a -> OffsetM b -> OffsetM c)
    -> (forall a b. OffsetM a -> OffsetM b -> OffsetM b)
    -> (forall a b. OffsetM a -> OffsetM b -> OffsetM a)
    -> Applicative OffsetM
    forall a. a -> OffsetM a
    forall a b. OffsetM a -> OffsetM b -> OffsetM a
    forall a b. OffsetM a -> OffsetM b -> OffsetM b
    forall a b. OffsetM (a -> b) -> OffsetM a -> OffsetM b
    forall a b c. (a -> b -> c) -> OffsetM a -> OffsetM b -> OffsetM c
    forall (f :: * -> *).
    Functor f =>
    (forall a. a -> f a)
    -> (forall a b. f (a -> b) -> f a -> f b)
    -> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
    -> (forall a b. f a -> f b -> f b)
    -> (forall a b. f a -> f b -> f a)
    -> Applicative f
    $cpure :: forall a. a -> OffsetM a
    pure :: forall a. a -> OffsetM a
    $c<*> :: forall a b. OffsetM (a -> b) -> OffsetM a -> OffsetM b
    <*> :: forall a b. OffsetM (a -> b) -> OffsetM a -> OffsetM b
    $cliftA2 :: forall a b c. (a -> b -> c) -> OffsetM a -> OffsetM b -> OffsetM c
    liftA2 :: forall a b c. (a -> b -> c) -> OffsetM a -> OffsetM b -> OffsetM c
    $c*> :: forall a b. OffsetM a -> OffsetM b -> OffsetM b
    *> :: forall a b. OffsetM a -> OffsetM b -> OffsetM b
    $c<* :: forall a b. OffsetM a -> OffsetM b -> OffsetM a
    <* :: forall a b. OffsetM a -> OffsetM b -> OffsetM a
    Applicative,
          (forall a b. (a -> b) -> OffsetM a -> OffsetM b)
    -> (forall a b. a -> OffsetM b -> OffsetM a) -> Functor OffsetM
    forall a b. a -> OffsetM b -> OffsetM a
    forall a b. (a -> b) -> OffsetM a -> OffsetM b
    forall (f :: * -> *).
    (forall a b. (a -> b) -> f a -> f b)
    -> (forall a b. a -> f b -> f a) -> Functor f
    $cfmap :: forall a b. (a -> b) -> OffsetM a -> OffsetM b
    fmap :: forall a b. (a -> b) -> OffsetM a -> OffsetM b
    $c<$ :: forall a b. a -> OffsetM b -> OffsetM a
    <$ :: forall a b. a -> OffsetM b -> OffsetM a
    Functor,
          Applicative OffsetM
    Applicative OffsetM =>
    (forall a b. OffsetM a -> (a -> OffsetM b) -> OffsetM b)
    -> (forall a b. OffsetM a -> OffsetM b -> OffsetM b)
    -> (forall a. a -> OffsetM a)
    -> Monad OffsetM
    forall a. a -> OffsetM a
    forall a b. OffsetM a -> OffsetM b -> OffsetM b
    forall a b. OffsetM a -> (a -> OffsetM b) -> OffsetM b
    forall (m :: * -> *).
    Applicative m =>
    (forall a b. m a -> (a -> m b) -> m b)
    -> (forall a b. m a -> m b -> m b)
    -> (forall a. a -> m a)
    -> Monad m
    $c>>= :: forall a b. OffsetM a -> (a -> OffsetM b) -> OffsetM b
    >>= :: forall a b. OffsetM a -> (a -> OffsetM b) -> OffsetM b
    $c>> :: forall a b. OffsetM a -> OffsetM b -> OffsetM b
    >> :: forall a b. OffsetM a -> OffsetM b -> OffsetM b
    $creturn :: forall a. a -> OffsetM a
    return :: forall a. a -> OffsetM a
    Monad,
          HasScope GPUMem,
          LocalScope GPUMem,
          MonadError String,
          Monad OffsetM
    OffsetM VNameSource
    Monad OffsetM =>
    OffsetM VNameSource
    -> (VNameSource -> OffsetM ()) -> MonadFreshNames OffsetM
    VNameSource -> OffsetM ()
    forall (m :: * -> *).
    Monad m =>
    m VNameSource -> (VNameSource -> m ()) -> MonadFreshNames m
    $cgetNameSource :: OffsetM VNameSource
    getNameSource :: OffsetM VNameSource
    $cputNameSource :: VNameSource -> OffsetM ()
    putNameSource :: VNameSource -> OffsetM ()
    MonadFreshNames
        )

    instance MonadBuilder OffsetM where
      type Rep OffsetM = GPUMem
      mkExpDecM :: Pat (LetDec (Rep OffsetM))
    -> Exp (Rep OffsetM) -> OffsetM (ExpDec (Rep OffsetM))
    mkExpDecM Pat (LetDec (Rep OffsetM))
    pat Exp (Rep OffsetM)
    e = BuilderT
      GPUMem (StateT VNameSource (Either String)) (ExpDec (Rep OffsetM))
    -> OffsetM (ExpDec (Rep OffsetM))
    forall a.
    BuilderT GPUMem (StateT VNameSource (Either String)) a -> OffsetM a
    OffsetM (BuilderT
       GPUMem (StateT VNameSource (Either String)) (ExpDec (Rep OffsetM))
     -> OffsetM (ExpDec (Rep OffsetM)))
    -> BuilderT
         GPUMem (StateT VNameSource (Either String)) (ExpDec (Rep OffsetM))
    -> OffsetM (ExpDec (Rep OffsetM))
    forall a b. (a -> b) -> a -> b
    $ Pat
      (LetDec
         (Rep (BuilderT GPUMem (StateT VNameSource (Either String)))))
    -> Exp (Rep (BuilderT GPUMem (StateT VNameSource (Either String))))
    -> BuilderT
         GPUMem
         (StateT VNameSource (Either String))
         (ExpDec
            (Rep (BuilderT GPUMem (StateT VNameSource (Either String)))))
    forall (m :: * -> *).
    MonadBuilder m =>
    Pat (LetDec (Rep m)) -> Exp (Rep m) -> m (ExpDec (Rep m))
    mkExpDecM Pat
      (LetDec
         (Rep (BuilderT GPUMem (StateT VNameSource (Either String)))))
    Pat (LetDec (Rep OffsetM))
    pat Exp (Rep (BuilderT GPUMem (StateT VNameSource (Either String))))
    Exp (Rep OffsetM)
    e
      mkBodyM :: Stms (Rep OffsetM) -> Result -> OffsetM (Body (Rep OffsetM))
    mkBodyM Stms (Rep OffsetM)
    stms Result
    res = BuilderT
      GPUMem (StateT VNameSource (Either String)) (Body (Rep OffsetM))
    -> OffsetM (Body (Rep OffsetM))
    forall a.
    BuilderT GPUMem (StateT VNameSource (Either String)) a -> OffsetM a
    OffsetM (BuilderT
       GPUMem (StateT VNameSource (Either String)) (Body (Rep OffsetM))
     -> OffsetM (Body (Rep OffsetM)))
    -> BuilderT
         GPUMem (StateT VNameSource (Either String)) (Body (Rep OffsetM))
    -> OffsetM (Body (Rep OffsetM))
    forall a b. (a -> b) -> a -> b
    $ Stms (Rep (BuilderT GPUMem (StateT VNameSource (Either String))))
    -> Result
    -> BuilderT
         GPUMem
         (StateT VNameSource (Either String))
         (Body (Rep (BuilderT GPUMem (StateT VNameSource (Either String)))))
    forall (m :: * -> *).
    MonadBuilder m =>
    Stms (Rep m) -> Result -> m (Body (Rep m))
    mkBodyM Stms (Rep (BuilderT GPUMem (StateT VNameSource (Either String))))
    Stms (Rep OffsetM)
    stms Result
    res
      mkLetNamesM :: [VName] -> Exp (Rep OffsetM) -> OffsetM (Stm (Rep OffsetM))
    mkLetNamesM [VName]
    pat Exp (Rep OffsetM)
    e = BuilderT
      GPUMem (StateT VNameSource (Either String)) (Stm (Rep OffsetM))
    -> OffsetM (Stm (Rep OffsetM))
    forall a.
    BuilderT GPUMem (StateT VNameSource (Either String)) a -> OffsetM a
    OffsetM (BuilderT
       GPUMem (StateT VNameSource (Either String)) (Stm (Rep OffsetM))
     -> OffsetM (Stm (Rep OffsetM)))
    -> BuilderT
         GPUMem (StateT VNameSource (Either String)) (Stm (Rep OffsetM))
    -> OffsetM (Stm (Rep OffsetM))
    forall a b. (a -> b) -> a -> b
    $ [VName]
    -> Exp (Rep (BuilderT GPUMem (StateT VNameSource (Either String))))
    -> BuilderT
         GPUMem
         (StateT VNameSource (Either String))
         (Stm (Rep (BuilderT GPUMem (StateT VNameSource (Either String)))))
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m (Stm (Rep m))
    mkLetNamesM [VName]
    pat Exp (Rep (BuilderT GPUMem (StateT VNameSource (Either String))))
    Exp (Rep OffsetM)
    e

      addStms :: Stms (Rep OffsetM) -> OffsetM ()
    addStms = BuilderT GPUMem (StateT VNameSource (Either String)) ()
    -> OffsetM ()
    forall a.
    BuilderT GPUMem (StateT VNameSource (Either String)) a -> OffsetM a
    OffsetM (BuilderT GPUMem (StateT VNameSource (Either String)) ()
     -> OffsetM ())
    -> (Stms GPUMem
        -> BuilderT GPUMem (StateT VNameSource (Either String)) ())
    -> Stms GPUMem
    -> OffsetM ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Stms (Rep (BuilderT GPUMem (StateT VNameSource (Either String))))
    -> BuilderT GPUMem (StateT VNameSource (Either String)) ()
    Stms GPUMem
    -> BuilderT GPUMem (StateT VNameSource (Either String)) ()
    forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
    addStms
      collectStms :: forall a. OffsetM a -> OffsetM (a, Stms (Rep OffsetM))
    collectStms (OffsetM BuilderT GPUMem (StateT VNameSource (Either String)) a
    m) = BuilderT
      GPUMem (StateT VNameSource (Either String)) (a, Stms (Rep OffsetM))
    -> OffsetM (a, Stms (Rep OffsetM))
    forall a.
    BuilderT GPUMem (StateT VNameSource (Either String)) a -> OffsetM a
    OffsetM (BuilderT
       GPUMem (StateT VNameSource (Either String)) (a, Stms (Rep OffsetM))
     -> OffsetM (a, Stms (Rep OffsetM)))
    -> BuilderT
         GPUMem (StateT VNameSource (Either String)) (a, Stms (Rep OffsetM))
    -> OffsetM (a, Stms (Rep OffsetM))
    forall a b. (a -> b) -> a -> b
    $ BuilderT GPUMem (StateT VNameSource (Either String)) a
    -> BuilderT
         GPUMem
         (StateT VNameSource (Either String))
         (a,
          Stms (Rep (BuilderT GPUMem (StateT VNameSource (Either String)))))
    forall a.
    BuilderT GPUMem (StateT VNameSource (Either String)) a
    -> BuilderT
         GPUMem
         (StateT VNameSource (Either String))
         (a,
          Stms (Rep (BuilderT GPUMem (StateT VNameSource (Either String)))))
    forall (m :: * -> *) a.
    MonadBuilder m =>
    m a -> m (a, Stms (Rep m))
    collectStms BuilderT GPUMem (StateT VNameSource (Either String)) a
    m

    runOffsetM ::
      (MonadFreshNames m) =>
      Scope GPUMem ->
      OffsetM a ->
      m (Either String a)
    runOffsetM :: forall (m :: * -> *) a.
    MonadFreshNames m =>
    Scope GPUMem -> OffsetM a -> m (Either String a)
    runOffsetM Scope GPUMem
    scope (OffsetM BuilderT GPUMem (StateT VNameSource (Either String)) a
    m) = (VNameSource -> (Either String a, VNameSource))
    -> m (Either String a)
    forall (m :: * -> *) a.
    MonadFreshNames m =>
    (VNameSource -> (a, VNameSource)) -> m a
    modifyNameSource ((VNameSource -> (Either String a, VNameSource))
     -> m (Either String a))
    -> (VNameSource -> (Either String a, VNameSource))
    -> m (Either String a)
    forall a b. (a -> b) -> a -> b
    $ \VNameSource
    src ->
      case StateT VNameSource (Either String) (a, Stms GPUMem)
    -> VNameSource -> Either String ((a, Stms GPUMem), VNameSource)
    forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
    runStateT (BuilderT GPUMem (StateT VNameSource (Either String)) a
    -> Scope GPUMem
    -> StateT VNameSource (Either String) (a, Stms GPUMem)
    forall (m :: * -> *) rep a.
    MonadFreshNames m =>
    BuilderT rep m a -> Scope rep -> m (a, Stms rep)
    runBuilderT BuilderT GPUMem (StateT VNameSource (Either String)) a
    m Scope GPUMem
    scope) VNameSource
    src of
        Left String
    e -> (String -> Either String a
    forall a b. a -> Either a b
    Left String
    e, VNameSource
    src)
        Right ((a, Stms GPUMem)
    x, VNameSource
    src') -> (a -> Either String a
    forall a b. b -> Either a b
    Right ((a, Stms GPUMem) -> a
    forall a b. (a, b) -> a
    fst (a, Stms GPUMem)
    x), VNameSource
    src')

    lookupNewBase :: VName -> [Exp64] -> RebaseMap -> Maybe Expansion
    lookupNewBase :: VName
    -> [TPrimExp Int64 VName]
    -> RebaseMap
    -> Maybe (TPrimExp Int64 VName, TPrimExp Int64 VName)
    lookupNewBase VName
    name [TPrimExp Int64 VName]
    x RebaseMap
    offsets =
      (([TPrimExp Int64 VName]
     -> (TPrimExp Int64 VName, TPrimExp Int64 VName))
    -> [TPrimExp Int64 VName]
    -> (TPrimExp Int64 VName, TPrimExp Int64 VName)
    forall a b. (a -> b) -> a -> b
    $ [TPrimExp Int64 VName]
    x) (([TPrimExp Int64 VName]
      -> (TPrimExp Int64 VName, TPrimExp Int64 VName))
     -> (TPrimExp Int64 VName, TPrimExp Int64 VName))
    -> Maybe
         ([TPrimExp Int64 VName]
          -> (TPrimExp Int64 VName, TPrimExp Int64 VName))
    -> Maybe (TPrimExp Int64 VName, TPrimExp Int64 VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> VName
    -> RebaseMap
    -> Maybe
         ([TPrimExp Int64 VName]
          -> (TPrimExp Int64 VName, TPrimExp Int64 VName))
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    name RebaseMap
    offsets

    offsetMemoryInKernelBody :: RebaseMap -> KernelBody GPUMem -> OffsetM (KernelBody GPUMem)
    offsetMemoryInKernelBody :: RebaseMap -> KernelBody GPUMem -> OffsetM (KernelBody GPUMem)
    offsetMemoryInKernelBody RebaseMap
    offsets KernelBody GPUMem
    kbody = do
      stms' <-
        OffsetM () -> OffsetM (Stms (Rep OffsetM))
    forall (m :: * -> *) a. MonadBuilder m => m a -> m (Stms (Rep m))
    collectStms_ (OffsetM () -> OffsetM (Stms (Rep OffsetM)))
    -> OffsetM () -> OffsetM (Stms (Rep OffsetM))
    forall a b. (a -> b) -> a -> b
    $
          (Stm GPUMem -> OffsetM ()) -> Stms GPUMem -> OffsetM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Stm (Rep OffsetM) -> OffsetM ()
    Stm GPUMem -> OffsetM ()
    forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
    addStm (Stm GPUMem -> OffsetM ())
    -> (Stm GPUMem -> OffsetM (Stm GPUMem)) -> Stm GPUMem -> OffsetM ()
    forall (m :: * -> *) b c a.
    Monad m =>
    (b -> m c) -> (a -> m b) -> a -> m c
    <=< RebaseMap -> Stm GPUMem -> OffsetM (Stm GPUMem)
    offsetMemoryInStm RebaseMap
    offsets) (KernelBody GPUMem -> Stms GPUMem
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody GPUMem
    kbody)
      pure kbody {kernelBodyStms = stms'}

    offsetMemoryInBody :: RebaseMap -> Body GPUMem -> OffsetM (Body GPUMem)
    offsetMemoryInBody :: RebaseMap -> Body GPUMem -> OffsetM (Body GPUMem)
    offsetMemoryInBody RebaseMap
    offsets (Body BodyDec GPUMem
    _ Stms GPUMem
    stms Result
    res) = do
      OffsetM Result -> OffsetM (Body (Rep OffsetM))
    forall (m :: * -> *).
    MonadBuilder m =>
    m Result -> m (Body (Rep m))
    buildBody_ (OffsetM Result -> OffsetM (Body (Rep OffsetM)))
    -> OffsetM Result -> OffsetM (Body (Rep OffsetM))
    forall a b. (a -> b) -> a -> b
    $ do
        (Stm GPUMem -> OffsetM ()) -> Stms GPUMem -> OffsetM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Stm (Rep OffsetM) -> OffsetM ()
    Stm GPUMem -> OffsetM ()
    forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
    addStm (Stm GPUMem -> OffsetM ())
    -> (Stm GPUMem -> OffsetM (Stm GPUMem)) -> Stm GPUMem -> OffsetM ()
    forall (m :: * -> *) b c a.
    Monad m =>
    (b -> m c) -> (a -> m b) -> a -> m c
    <=< RebaseMap -> Stm GPUMem -> OffsetM (Stm GPUMem)
    offsetMemoryInStm RebaseMap
    offsets) Stms GPUMem
    stms
        Result -> OffsetM Result
    forall a. a -> OffsetM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Result
    res

    argsContext :: [SubExp] -> OffsetM [SubExp]
    argsContext :: [SubExp] -> OffsetM [SubExp]
    argsContext = ([[SubExp]] -> [SubExp]) -> OffsetM [[SubExp]] -> OffsetM [SubExp]
    forall a b. (a -> b) -> OffsetM a -> OffsetM b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [[SubExp]] -> [SubExp]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat (OffsetM [[SubExp]] -> OffsetM [SubExp])
    -> ([SubExp] -> OffsetM [[SubExp]]) -> [SubExp] -> OffsetM [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (SubExp -> OffsetM [SubExp]) -> [SubExp] -> OffsetM [[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 SubExp -> OffsetM [SubExp]
    forall {m :: * -> *} {inner :: * -> *}.
    (RetType (Rep m) ~ RetTypeMem,
     LParamInfo (Rep m) ~ MemInfo SubExp NoUniqueness MemBind,
     BranchType (Rep m) ~ BranchTypeMem,
     FParamInfo (Rep m) ~ MemInfo SubExp Uniqueness MemBind,
     OpC (Rep m) ~ MemOp inner, OpReturns inner, MonadBuilder m,
     RephraseOp inner, HasLetDecMem (LetDec (Rep m)),
     Pretty (inner (Rep m)), Rename (inner (Rep m)),
     Substitute (inner (Rep m)), FreeIn (inner (Rep m)),
     Show (inner (Rep m)), Ord (inner (Rep m))) =>
    SubExp -> m [SubExp]
    resCtx
      where
        resCtx :: SubExp -> m [SubExp]
    resCtx SubExp
    se = do
          v_t <- SubExp -> m (MemInfo SubExp NoUniqueness MemBind)
    forall rep (m :: * -> *) (inner :: * -> *).
    (HasScope rep m, Mem rep inner) =>
    SubExp -> m (MemInfo SubExp NoUniqueness MemBind)
    subExpMemInfo SubExp
    se
          case v_t of
            MemArray PrimType
    _ ShapeBase SubExp
    _ NoUniqueness
    _ (ArrayIn VName
    mem LMAD
    lmad) -> do
              ctxs <- (TPrimExp Int64 VName -> m SubExp)
    -> [TPrimExp Int64 VName] -> m [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 (String -> Exp (Rep m) -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    "ctx" (Exp (Rep m) -> m SubExp)
    -> (TPrimExp Int64 VName -> m (Exp (Rep m)))
    -> TPrimExp Int64 VName
    -> m SubExp
    forall (m :: * -> *) b c a.
    Monad m =>
    (b -> m c) -> (a -> m b) -> a -> m c
    <=< TPrimExp Int64 VName -> m (Exp (Rep m))
    forall a (m :: * -> *).
    (ToExp a, MonadBuilder m) =>
    a -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    TPrimExp Int64 VName -> m (Exp (Rep m))
    toExp) (LMAD -> [TPrimExp Int64 VName]
    forall a. LMAD a -> [a]
    LMAD.existentialized LMAD
    lmad)
              pure $ Var mem : ctxs
            MemInfo SubExp NoUniqueness MemBind
    _ -> [SubExp] -> m [SubExp]
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure []

    offsetMemoryInBodyReturnCtx :: RebaseMap -> Body GPUMem -> OffsetM (Body GPUMem)
    offsetMemoryInBodyReturnCtx :: RebaseMap -> Body GPUMem -> OffsetM (Body GPUMem)
    offsetMemoryInBodyReturnCtx RebaseMap
    offsets (Body BodyDec GPUMem
    _ Stms GPUMem
    stms Result
    res) = do
      OffsetM Result -> OffsetM (Body (Rep OffsetM))
    forall (m :: * -> *).
    MonadBuilder m =>
    m Result -> m (Body (Rep m))
    buildBody_ (OffsetM Result -> OffsetM (Body (Rep OffsetM)))
    -> OffsetM Result -> OffsetM (Body (Rep OffsetM))
    forall a b. (a -> b) -> a -> b
    $ do
        (Stm GPUMem -> OffsetM ()) -> Stms GPUMem -> OffsetM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Stm (Rep OffsetM) -> OffsetM ()
    Stm GPUMem -> OffsetM ()
    forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
    addStm (Stm GPUMem -> OffsetM ())
    -> (Stm GPUMem -> OffsetM (Stm GPUMem)) -> Stm GPUMem -> OffsetM ()
    forall (m :: * -> *) b c a.
    Monad m =>
    (b -> m c) -> (a -> m b) -> a -> m c
    <=< RebaseMap -> Stm GPUMem -> OffsetM (Stm GPUMem)
    offsetMemoryInStm RebaseMap
    offsets) Stms GPUMem
    stms
        ctx <- [SubExp] -> OffsetM [SubExp]
    argsContext ([SubExp] -> OffsetM [SubExp]) -> [SubExp] -> OffsetM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ (SubExpRes -> SubExp) -> Result -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExpRes -> SubExp
    resSubExp Result
    res
        pure $ res <> subExpsRes ctx

    lmadFrom :: LMAD.Shape num -> [num] -> LMAD.LMAD num
    lmadFrom :: forall num. Shape num -> Shape num -> LMAD num
    lmadFrom Shape num
    shape Shape num
    xs =
      num -> [LMADDim num] -> LMAD num
    forall num. num -> [LMADDim num] -> LMAD num
    LMAD.LMAD (Shape num -> num
    forall a. HasCallStack => [a] -> a
    head Shape num
    xs) ([LMADDim num] -> LMAD num) -> [LMADDim num] -> LMAD num
    forall a b. (a -> b) -> a -> b
    $ (num -> num -> LMADDim num)
    -> Shape num -> Shape num -> [LMADDim num]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith num -> num -> LMADDim num
    forall num. num -> num -> LMADDim num
    LMAD.LMADDim (Int -> Shape num -> Shape num
    forall a. Int -> [a] -> [a]
    drop Int
    1 Shape num
    xs) Shape num
    shape

    -- | Append pattern elements corresponding to memory and index
    -- function components for every array bound in the pattern.
    addPatternContext :: Pat LetDecMem -> OffsetM (Pat LetDecMem)
    addPatternContext :: Pat (MemInfo SubExp NoUniqueness MemBind)
    -> OffsetM (Pat (MemInfo SubExp NoUniqueness MemBind))
    addPatternContext (Pat [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    pes) = Scope GPUMem
    -> OffsetM (Pat (MemInfo SubExp NoUniqueness MemBind))
    -> OffsetM (Pat (MemInfo SubExp NoUniqueness MemBind))
    forall a. Scope GPUMem -> OffsetM a -> OffsetM a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope (Pat (MemInfo SubExp NoUniqueness MemBind) -> Scope GPUMem
    forall rep dec. (LetDec rep ~ dec) => Pat dec -> Scope rep
    scopeOfPat ([PatElem (MemInfo SubExp NoUniqueness MemBind)]
    -> Pat (MemInfo SubExp NoUniqueness MemBind)
    forall dec. [PatElem dec] -> Pat dec
    Pat [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    pes)) (OffsetM (Pat (MemInfo SubExp NoUniqueness MemBind))
     -> OffsetM (Pat (MemInfo SubExp NoUniqueness MemBind)))
    -> OffsetM (Pat (MemInfo SubExp NoUniqueness MemBind))
    -> OffsetM (Pat (MemInfo SubExp NoUniqueness MemBind))
    forall a b. (a -> b) -> a -> b
    $ do
      (pes_ctx, pes') <- ([PatElem (MemInfo SubExp NoUniqueness MemBind)]
     -> PatElem (MemInfo SubExp NoUniqueness MemBind)
     -> OffsetM
          ([PatElem (MemInfo SubExp NoUniqueness MemBind)],
           PatElem (MemInfo SubExp NoUniqueness MemBind)))
    -> [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    -> [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    -> OffsetM
         ([PatElem (MemInfo SubExp NoUniqueness MemBind)],
          [PatElem (MemInfo SubExp NoUniqueness MemBind)])
    forall (m :: * -> *) (t :: * -> *) acc x y.
    (Monad m, Traversable t) =>
    (acc -> x -> m (acc, y)) -> acc -> t x -> m (acc, t y)
    mapAccumLM [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    -> PatElem (MemInfo SubExp NoUniqueness MemBind)
    -> OffsetM
         ([PatElem (MemInfo SubExp NoUniqueness MemBind)],
          PatElem (MemInfo SubExp NoUniqueness MemBind))
    forall {rep} {inner :: * -> *} {m :: * -> *} {d} {u} {ret} {d} {u}.
    (BranchType rep ~ BranchTypeMem,
     LParamInfo rep ~ MemInfo SubExp NoUniqueness MemBind,
     FParamInfo rep ~ MemInfo SubExp Uniqueness MemBind,
     RetType rep ~ RetTypeMem, OpC rep ~ MemOp inner,
     HasLetDecMem (LetDec rep), ASTRep rep, OpReturns inner,
     RephraseOp inner, HasScope rep m, MonadFreshNames m,
     Pretty (inner rep), Rename (inner rep), Show (inner rep),
     Ord (inner rep), Substitute (inner rep), FreeIn (inner rep)) =>
    [PatElem (MemInfo d u ret)]
    -> PatElem (MemInfo d u MemBind)
    -> m ([PatElem (MemInfo d u ret)], PatElem (MemInfo d u MemBind))
    onType [] [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    pes
      pure $ Pat $ pes' <> pes_ctx
      where
        onType :: [PatElem (MemInfo d u ret)]
    -> PatElem (MemInfo d u MemBind)
    -> m ([PatElem (MemInfo d u ret)], PatElem (MemInfo d u MemBind))
    onType
          [PatElem (MemInfo d u ret)]
    acc
          (PatElem VName
    pe_v (MemArray PrimType
    pt ShapeBase d
    pe_shape u
    pe_u (ArrayIn VName
    pe_mem LMAD
    lmad))) = do
            space <- VName -> m Space
    forall rep (inner :: * -> *) (m :: * -> *).
    (Mem rep inner, HasScope rep m, Monad m) =>
    VName -> m Space
    lookupMemSpace VName
    pe_mem
            pe_mem' <- newVName $ baseString pe_mem <> "_ext"
            let num_exts = [TPrimExp Int64 VName] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length (LMAD -> [TPrimExp Int64 VName]
    forall a. LMAD a -> [a]
    LMAD.existentialized LMAD
    lmad)
            lmad_exts <-
              replicateM num_exts $
                PatElem <$> newVName "ext" <*> pure (MemPrim int64)
            let pe_lmad' = [TPrimExp Int64 VName] -> [TPrimExp Int64 VName] -> LMAD
    forall num. Shape num -> Shape num -> LMAD num
    lmadFrom (LMAD -> [TPrimExp Int64 VName]
    forall a. LMAD a -> [a]
    LMAD.shape LMAD
    lmad) ([TPrimExp Int64 VName] -> LMAD) -> [TPrimExp Int64 VName] -> LMAD
    forall a b. (a -> b) -> a -> b
    $ (PatElem (MemInfo d u ret) -> TPrimExp Int64 VName)
    -> [PatElem (MemInfo d u ret)] -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 (VName -> TPrimExp Int64 VName)
    -> (PatElem (MemInfo d u ret) -> VName)
    -> PatElem (MemInfo d u ret)
    -> TPrimExp Int64 VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . PatElem (MemInfo d u ret) -> VName
    forall dec. PatElem dec -> VName
    patElemName) [PatElem (MemInfo d u ret)]
    lmad_exts
            pure
              ( acc ++ PatElem pe_mem' (MemMem space) : lmad_exts,
                PatElem pe_v $ MemArray pt pe_shape pe_u $ ArrayIn pe_mem' pe_lmad'
              )
        onType [PatElem (MemInfo d u ret)]
    acc PatElem (MemInfo d u MemBind)
    t = ([PatElem (MemInfo d u ret)], PatElem (MemInfo d u MemBind))
    -> m ([PatElem (MemInfo d u ret)], PatElem (MemInfo d u MemBind))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([PatElem (MemInfo d u ret)]
    acc, PatElem (MemInfo d u MemBind)
    t)

    -- | Append pattern elements corresponding to memory and index
    -- function components for every array bound in the parameters.
    addParamsContext :: [Param FParamMem] -> OffsetM [Param FParamMem]
    addParamsContext :: [Param (MemInfo SubExp Uniqueness MemBind)]
    -> OffsetM [Param (MemInfo SubExp Uniqueness MemBind)]
    addParamsContext [Param (MemInfo SubExp Uniqueness MemBind)]
    ps = Scope GPUMem
    -> OffsetM [Param (MemInfo SubExp Uniqueness MemBind)]
    -> OffsetM [Param (MemInfo SubExp Uniqueness MemBind)]
    forall a. Scope GPUMem -> OffsetM a -> OffsetM a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope ([Param (MemInfo SubExp Uniqueness MemBind)] -> Scope GPUMem
    forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfFParams [Param (MemInfo SubExp Uniqueness MemBind)]
    ps) (OffsetM [Param (MemInfo SubExp Uniqueness MemBind)]
     -> OffsetM [Param (MemInfo SubExp Uniqueness MemBind)])
    -> OffsetM [Param (MemInfo SubExp Uniqueness MemBind)]
    -> OffsetM [Param (MemInfo SubExp Uniqueness MemBind)]
    forall a b. (a -> b) -> a -> b
    $ do
      (ps_ctx, ps') <- ([Param (MemInfo SubExp Uniqueness MemBind)]
     -> Param (MemInfo SubExp Uniqueness MemBind)
     -> OffsetM
          ([Param (MemInfo SubExp Uniqueness MemBind)],
           Param (MemInfo SubExp Uniqueness MemBind)))
    -> [Param (MemInfo SubExp Uniqueness MemBind)]
    -> [Param (MemInfo SubExp Uniqueness MemBind)]
    -> OffsetM
         ([Param (MemInfo SubExp Uniqueness MemBind)],
          [Param (MemInfo SubExp Uniqueness MemBind)])
    forall (m :: * -> *) (t :: * -> *) acc x y.
    (Monad m, Traversable t) =>
    (acc -> x -> m (acc, y)) -> acc -> t x -> m (acc, t y)
    mapAccumLM [Param (MemInfo SubExp Uniqueness MemBind)]
    -> Param (MemInfo SubExp Uniqueness MemBind)
    -> OffsetM
         ([Param (MemInfo SubExp Uniqueness MemBind)],
          Param (MemInfo SubExp Uniqueness MemBind))
    forall {rep} {inner :: * -> *} {m :: * -> *} {d} {u} {ret} {d} {u}.
    (BranchType rep ~ BranchTypeMem,
     LParamInfo rep ~ MemInfo SubExp NoUniqueness MemBind,
     FParamInfo rep ~ MemInfo SubExp Uniqueness MemBind,
     RetType rep ~ RetTypeMem, OpC rep ~ MemOp inner,
     HasLetDecMem (LetDec rep), ASTRep rep, OpReturns inner,
     RephraseOp inner, HasScope rep m, MonadFreshNames m,
     Pretty (inner rep), Rename (inner rep), Show (inner rep),
     Ord (inner rep), Substitute (inner rep), FreeIn (inner rep)) =>
    [Param (MemInfo d u ret)]
    -> Param (MemInfo d u MemBind)
    -> m ([Param (MemInfo d u ret)], Param (MemInfo d u MemBind))
    onType [] [Param (MemInfo SubExp Uniqueness MemBind)]
    ps
      pure $ ps' <> ps_ctx
      where
        onType :: [Param (MemInfo d u ret)]
    -> Param (MemInfo d u MemBind)
    -> m ([Param (MemInfo d u ret)], Param (MemInfo d u MemBind))
    onType [Param (MemInfo d u ret)]
    acc (Param Attrs
    attr VName
    v (MemArray PrimType
    pt ShapeBase d
    shape u
    u (ArrayIn VName
    mem LMAD
    lmad))) = do
          space <- VName -> m Space
    forall rep (inner :: * -> *) (m :: * -> *).
    (Mem rep inner, HasScope rep m, Monad m) =>
    VName -> m Space
    lookupMemSpace VName
    mem
          mem' <- newVName $ baseString mem <> "_ext"
          let num_exts = [TPrimExp Int64 VName] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length (LMAD -> [TPrimExp Int64 VName]
    forall a. LMAD a -> [a]
    LMAD.existentialized LMAD
    lmad)
          lmad_exts <-
            replicateM num_exts $
              Param mempty <$> newVName "ext" <*> pure (MemPrim int64)
          let lmad' = [TPrimExp Int64 VName] -> [TPrimExp Int64 VName] -> LMAD
    forall num. Shape num -> Shape num -> LMAD num
    lmadFrom (LMAD -> [TPrimExp Int64 VName]
    forall a. LMAD a -> [a]
    LMAD.shape LMAD
    lmad) ([TPrimExp Int64 VName] -> LMAD) -> [TPrimExp Int64 VName] -> LMAD
    forall a b. (a -> b) -> a -> b
    $ (Param (MemInfo d u ret) -> TPrimExp Int64 VName)
    -> [Param (MemInfo d u ret)] -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 (VName -> TPrimExp Int64 VName)
    -> (Param (MemInfo d u ret) -> VName)
    -> Param (MemInfo d u ret)
    -> TPrimExp Int64 VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Param (MemInfo d u ret) -> VName
    forall dec. Param dec -> VName
    paramName) [Param (MemInfo d u ret)]
    lmad_exts
          pure
            ( acc ++ Param mempty mem' (MemMem space) : lmad_exts,
              Param attr v $ MemArray pt shape u $ ArrayIn mem' lmad'
            )
        onType [Param (MemInfo d u ret)]
    acc Param (MemInfo d u MemBind)
    t = ([Param (MemInfo d u ret)], Param (MemInfo d u MemBind))
    -> m ([Param (MemInfo d u ret)], Param (MemInfo d u MemBind))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([Param (MemInfo d u ret)]
    acc, Param (MemInfo d u MemBind)
    t)

    offsetBranch ::
      Pat LetDecMem ->
      [BranchTypeMem] ->
      OffsetM (Pat LetDecMem, [BranchTypeMem])
    offsetBranch :: Pat (MemInfo SubExp NoUniqueness MemBind)
    -> [BranchTypeMem]
    -> OffsetM
         (Pat (MemInfo SubExp NoUniqueness MemBind), [BranchTypeMem])
    offsetBranch (Pat [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    pes) [BranchTypeMem]
    ts = do
      ((pes_ctx, ts_ctx), (pes', ts')) <-
        ([(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)]
     -> ([PatElem (MemInfo SubExp NoUniqueness MemBind)],
         [BranchTypeMem]))
    -> ([(PatElem (MemInfo SubExp NoUniqueness MemBind),
          BranchTypeMem)]
        -> ([PatElem (MemInfo SubExp NoUniqueness MemBind)],
            [BranchTypeMem]))
    -> ([(PatElem (MemInfo SubExp NoUniqueness MemBind),
          BranchTypeMem)],
        [(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)])
    -> (([PatElem (MemInfo SubExp NoUniqueness MemBind)],
         [BranchTypeMem]),
        ([PatElem (MemInfo SubExp NoUniqueness MemBind)], [BranchTypeMem]))
    forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
    forall (p :: * -> * -> *) a b c d.
    Bifunctor p =>
    (a -> b) -> (c -> d) -> p a c -> p b d
    bimap [(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)]
    -> ([PatElem (MemInfo SubExp NoUniqueness MemBind)],
        [BranchTypeMem])
    forall a b. [(a, b)] -> ([a], [b])
    unzip [(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)]
    -> ([PatElem (MemInfo SubExp NoUniqueness MemBind)],
        [BranchTypeMem])
    forall a b. [(a, b)] -> ([a], [b])
    unzip (([(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)],
      [(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)])
     -> (([PatElem (MemInfo SubExp NoUniqueness MemBind)],
          [BranchTypeMem]),
         ([PatElem (MemInfo SubExp NoUniqueness MemBind)],
          [BranchTypeMem])))
    -> OffsetM
         ([(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)],
          [(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)])
    -> OffsetM
         (([PatElem (MemInfo SubExp NoUniqueness MemBind)],
           [BranchTypeMem]),
          ([PatElem (MemInfo SubExp NoUniqueness MemBind)], [BranchTypeMem]))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> ([(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)]
     -> (PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)
     -> OffsetM
          ([(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)],
           (PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)))
    -> [(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)]
    -> [(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)]
    -> OffsetM
         ([(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)],
          [(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)])
    forall (m :: * -> *) (t :: * -> *) acc x y.
    (Monad m, Traversable t) =>
    (acc -> x -> m (acc, y)) -> acc -> t x -> m (acc, t y)
    mapAccumLM [(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)]
    -> (PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)
    -> OffsetM
         ([(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)],
          (PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem))
    onType [] ([PatElem (MemInfo SubExp NoUniqueness MemBind)]
    -> [BranchTypeMem]
    -> [(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    pes [BranchTypeMem]
    ts)
      pure (Pat $ pes' <> pes_ctx, ts' <> ts_ctx)
      where
        onType :: [(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)]
    -> (PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)
    -> OffsetM
         ([(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)],
          (PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem))
    onType
          [(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)]
    acc
          ( PatElem VName
    pe_v (MemArray PrimType
    _ ShapeBase SubExp
    pe_shape NoUniqueness
    pe_u (ArrayIn VName
    pe_mem LMAD
    pe_lmad)),
            MemArray PrimType
    pt ShapeBase ExtSize
    shape NoUniqueness
    u MemReturn
    meminfo
            ) = do
            (space, lmad) <- case MemReturn
    meminfo of
              ReturnsInBlock VName
    mem ExtLMAD
    lmad -> do
                space <- VName -> OffsetM Space
    forall rep (inner :: * -> *) (m :: * -> *).
    (Mem rep inner, HasScope rep m, Monad m) =>
    VName -> m Space
    lookupMemSpace VName
    mem
                pure (space, lmad)
              ReturnsNewBlock Space
    space Int
    _ ExtLMAD
    lmad ->
                (Space, ExtLMAD) -> OffsetM (Space, ExtLMAD)
    forall a. a -> OffsetM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Space
    space, ExtLMAD
    lmad)
            pe_mem' <- newVName $ baseString pe_mem <> "_ext"
            let start = [BranchTypeMem] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [BranchTypeMem]
    ts Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + [(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)]
    -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)]
    acc
                num_exts = [TPrimExp Int64 (Ext VName)] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length (ExtLMAD -> [TPrimExp Int64 (Ext VName)]
    forall a. LMAD a -> [a]
    LMAD.existentialized ExtLMAD
    lmad)
                ext (Free SubExp
    se) = VName -> Ext VName
    forall a. a -> Ext a
    Free (VName -> Ext VName)
    -> TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    se
                ext (Ext Int
    i) = Ext VName -> TPrimExp Int64 (Ext VName)
    forall a. a -> TPrimExp Int64 a
    le64 (Int -> Ext VName
    forall a. Int -> Ext a
    Ext Int
    i)
            lmad_exts <-
              replicateM num_exts $
                PatElem <$> newVName "ext" <*> pure (MemPrim int64)
            let pe_lmad' = [TPrimExp Int64 VName] -> [TPrimExp Int64 VName] -> LMAD
    forall num. Shape num -> Shape num -> LMAD num
    lmadFrom (LMAD -> [TPrimExp Int64 VName]
    forall a. LMAD a -> [a]
    LMAD.shape LMAD
    pe_lmad) ([TPrimExp Int64 VName] -> LMAD) -> [TPrimExp Int64 VName] -> LMAD
    forall a b. (a -> b) -> a -> b
    $ (PatElem (MemInfo SubExp NoUniqueness MemBind)
     -> TPrimExp Int64 VName)
    -> [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 (VName -> TPrimExp Int64 VName)
    -> (PatElem (MemInfo SubExp NoUniqueness MemBind) -> VName)
    -> PatElem (MemInfo SubExp NoUniqueness MemBind)
    -> TPrimExp Int64 VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . PatElem (MemInfo SubExp NoUniqueness MemBind) -> VName
    forall dec. PatElem dec -> VName
    patElemName) [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    lmad_exts
            pure
              ( acc
                  ++ (PatElem pe_mem' $ MemMem space, MemMem space)
                  : map (,MemPrim int64) lmad_exts,
                ( PatElem pe_v $ MemArray pt pe_shape pe_u $ ArrayIn pe_mem' pe_lmad',
                  MemArray pt shape u . ReturnsNewBlock space start . fmap ext $
                    LMAD.mkExistential (shapeDims shape) (1 + start)
                )
              )
        onType [(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)]
    acc (PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)
    t = ([(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)],
     (PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem))
    -> OffsetM
         ([(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)],
          (PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem))
    forall a. a -> OffsetM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([(PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)]
    acc, (PatElem (MemInfo SubExp NoUniqueness MemBind), BranchTypeMem)
    t)

    offsetMemoryInPat :: RebaseMap -> Pat LetDecMem -> [ExpReturns] -> Pat LetDecMem
    offsetMemoryInPat :: RebaseMap
    -> Pat (MemInfo SubExp NoUniqueness MemBind)
    -> [ExpReturns]
    -> Pat (MemInfo SubExp NoUniqueness MemBind)
    offsetMemoryInPat RebaseMap
    offsets (Pat [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    pes) [ExpReturns]
    rets = do
      [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    -> Pat (MemInfo SubExp NoUniqueness MemBind)
    forall dec. [PatElem dec] -> Pat dec
    Pat ([PatElem (MemInfo SubExp NoUniqueness MemBind)]
     -> Pat (MemInfo SubExp NoUniqueness MemBind))
    -> [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    -> Pat (MemInfo SubExp NoUniqueness MemBind)
    forall a b. (a -> b) -> a -> b
    $ (PatElem (MemInfo SubExp NoUniqueness MemBind)
     -> ExpReturns -> PatElem (MemInfo SubExp NoUniqueness MemBind))
    -> [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    -> [ExpReturns]
    -> [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith PatElem (MemInfo SubExp NoUniqueness MemBind)
    -> ExpReturns -> PatElem (MemInfo SubExp NoUniqueness MemBind)
    onPE [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    pes [ExpReturns]
    rets
      where
        onPE :: PatElem (MemInfo SubExp NoUniqueness MemBind)
    -> ExpReturns -> PatElem (MemInfo SubExp NoUniqueness MemBind)
    onPE
          (PatElem VName
    name (MemArray PrimType
    pt ShapeBase SubExp
    shape NoUniqueness
    u (ArrayIn VName
    mem LMAD
    _)))
          (MemArray PrimType
    _ ShapeBase ExtSize
    _ NoUniqueness
    _ Maybe MemReturn
    info)
            | Just ExtLMAD
    lmad <- Maybe MemReturn -> Maybe ExtLMAD
    getLMAD Maybe MemReturn
    info =
                VName
    -> MemInfo SubExp NoUniqueness MemBind
    -> PatElem (MemInfo SubExp NoUniqueness MemBind)
    forall dec. VName -> dec -> PatElem dec
    PatElem VName
    name (MemInfo SubExp NoUniqueness MemBind
     -> PatElem (MemInfo SubExp NoUniqueness MemBind))
    -> (LMAD -> MemInfo SubExp NoUniqueness MemBind)
    -> LMAD
    -> PatElem (MemInfo SubExp NoUniqueness MemBind)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . PrimType
    -> ShapeBase SubExp
    -> NoUniqueness
    -> MemBind
    -> MemInfo SubExp NoUniqueness MemBind
    forall d u ret.
    PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
    MemArray PrimType
    pt ShapeBase SubExp
    shape NoUniqueness
    u (MemBind -> MemInfo SubExp NoUniqueness MemBind)
    -> (LMAD -> MemBind) -> LMAD -> MemInfo SubExp NoUniqueness MemBind
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> LMAD -> MemBind
    ArrayIn VName
    mem (LMAD -> PatElem (MemInfo SubExp NoUniqueness MemBind))
    -> LMAD -> PatElem (MemInfo SubExp NoUniqueness MemBind)
    forall a b. (a -> b) -> a -> b
    $
                  (TPrimExp Int64 (Ext VName) -> TPrimExp Int64 VName)
    -> ExtLMAD -> LMAD
    forall a b. (a -> b) -> LMAD a -> LMAD b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap ((Ext VName -> VName)
    -> TPrimExp Int64 (Ext VName) -> TPrimExp Int64 VName
    forall a b. (a -> b) -> TPrimExp Int64 a -> TPrimExp Int64 b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap Ext VName -> VName
    unExt) ExtLMAD
    lmad
        onPE PatElem (MemInfo SubExp NoUniqueness MemBind)
    pe ExpReturns
    _ =
          RebaseMap
    -> MemInfo SubExp NoUniqueness MemBind
    -> MemInfo SubExp NoUniqueness MemBind
    forall u. RebaseMap -> MemBound u -> MemBound u
    offsetMemoryInMemBound RebaseMap
    offsets (MemInfo SubExp NoUniqueness MemBind
     -> MemInfo SubExp NoUniqueness MemBind)
    -> PatElem (MemInfo SubExp NoUniqueness MemBind)
    -> PatElem (MemInfo SubExp NoUniqueness MemBind)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> PatElem (MemInfo SubExp NoUniqueness MemBind)
    pe
        unExt :: Ext VName -> VName
    unExt (Ext Int
    i) = PatElem (MemInfo SubExp NoUniqueness MemBind) -> VName
    forall dec. PatElem dec -> VName
    patElemName ([PatElem (MemInfo SubExp NoUniqueness MemBind)]
    pes [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    -> Int -> PatElem (MemInfo SubExp NoUniqueness MemBind)
    forall a. HasCallStack => [a] -> Int -> a
    !! Int
    i)
        unExt (Free VName
    v) = VName
    v
        getLMAD :: Maybe MemReturn -> Maybe ExtLMAD
    getLMAD (Just (ReturnsNewBlock Space
    _ Int
    _ ExtLMAD
    lmad)) = ExtLMAD -> Maybe ExtLMAD
    forall a. a -> Maybe a
    Just ExtLMAD
    lmad
        getLMAD (Just (ReturnsInBlock VName
    _ ExtLMAD
    lmad)) = ExtLMAD -> Maybe ExtLMAD
    forall a. a -> Maybe a
    Just ExtLMAD
    lmad
        getLMAD Maybe MemReturn
    _ = Maybe ExtLMAD
    forall a. Maybe a
    Nothing

    offsetMemoryInParam :: RebaseMap -> Param (MemBound u) -> Param (MemBound u)
    offsetMemoryInParam :: forall u. RebaseMap -> Param (MemBound u) -> Param (MemBound u)
    offsetMemoryInParam RebaseMap
    offsets = (MemBound u -> MemBound u)
    -> Param (MemBound u) -> Param (MemBound u)
    forall a b. (a -> b) -> Param a -> Param b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap ((MemBound u -> MemBound u)
     -> Param (MemBound u) -> Param (MemBound u))
    -> (MemBound u -> MemBound u)
    -> Param (MemBound u)
    -> Param (MemBound u)
    forall a b. (a -> b) -> a -> b
    $ RebaseMap -> MemBound u -> MemBound u
    forall u. RebaseMap -> MemBound u -> MemBound u
    offsetMemoryInMemBound RebaseMap
    offsets

    offsetMemoryInMemBound :: RebaseMap -> MemBound u -> MemBound u
    offsetMemoryInMemBound :: forall u. RebaseMap -> MemBound u -> MemBound u
    offsetMemoryInMemBound RebaseMap
    offsets (MemArray PrimType
    pt ShapeBase SubExp
    shape u
    u (ArrayIn VName
    mem LMAD
    lmad))
      | Just (TPrimExp Int64 VName
    o, TPrimExp Int64 VName
    p) <- VName
    -> [TPrimExp Int64 VName]
    -> RebaseMap
    -> Maybe (TPrimExp Int64 VName, TPrimExp Int64 VName)
    lookupNewBase VName
    mem (LMAD -> [TPrimExp Int64 VName]
    forall a. LMAD a -> [a]
    LMAD.shape LMAD
    lmad) RebaseMap
    offsets =
          PrimType
    -> ShapeBase SubExp -> u -> MemBind -> MemInfo SubExp u MemBind
    forall d u ret.
    PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
    MemArray PrimType
    pt ShapeBase SubExp
    shape u
    u (MemBind -> MemInfo SubExp u MemBind)
    -> MemBind -> MemInfo SubExp u MemBind
    forall a b. (a -> b) -> a -> b
    $ VName -> LMAD -> MemBind
    ArrayIn VName
    mem (LMAD -> MemBind) -> LMAD -> MemBind
    forall a b. (a -> b) -> a -> b
    $ TPrimExp Int64 VName -> TPrimExp Int64 VName -> LMAD -> LMAD
    forall num. IntegralExp num => num -> num -> LMAD num -> LMAD num
    LMAD.expand TPrimExp Int64 VName
    o TPrimExp Int64 VName
    p LMAD
    lmad
    offsetMemoryInMemBound RebaseMap
    _ MemInfo SubExp u MemBind
    info = MemInfo SubExp u MemBind
    info

    offsetMemoryInBodyReturns :: RebaseMap -> BodyReturns -> BodyReturns
    offsetMemoryInBodyReturns :: RebaseMap -> BranchTypeMem -> BranchTypeMem
    offsetMemoryInBodyReturns RebaseMap
    offsets (MemArray PrimType
    pt ShapeBase ExtSize
    shape NoUniqueness
    u (ReturnsInBlock VName
    mem ExtLMAD
    lmad))
      | Just LMAD
    lmad' <- ExtLMAD -> Maybe LMAD
    isStaticLMAD ExtLMAD
    lmad,
        Just (TPrimExp Int64 VName
    o, TPrimExp Int64 VName
    p) <- VName
    -> [TPrimExp Int64 VName]
    -> RebaseMap
    -> Maybe (TPrimExp Int64 VName, TPrimExp Int64 VName)
    lookupNewBase VName
    mem (LMAD -> [TPrimExp Int64 VName]
    forall a. LMAD a -> [a]
    LMAD.shape LMAD
    lmad') RebaseMap
    offsets =
          PrimType
    -> ShapeBase ExtSize -> NoUniqueness -> MemReturn -> BranchTypeMem
    forall d u ret.
    PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
    MemArray PrimType
    pt ShapeBase ExtSize
    shape NoUniqueness
    u (MemReturn -> BranchTypeMem) -> MemReturn -> BranchTypeMem
    forall a b. (a -> b) -> a -> b
    $
            VName -> ExtLMAD -> MemReturn
    ReturnsInBlock VName
    mem (ExtLMAD -> MemReturn) -> ExtLMAD -> MemReturn
    forall a b. (a -> b) -> a -> b
    $
              TPrimExp Int64 (Ext VName)
    -> TPrimExp Int64 (Ext VName) -> ExtLMAD -> ExtLMAD
    forall num. IntegralExp num => num -> num -> LMAD num -> LMAD num
    LMAD.expand (VName -> Ext VName
    forall a. a -> Ext a
    Free (VName -> Ext VName)
    -> TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> TPrimExp Int64 VName
    o) ((VName -> Ext VName)
    -> TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName)
    forall a b. (a -> b) -> TPrimExp Int64 a -> TPrimExp Int64 b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap VName -> Ext VName
    forall a. a -> Ext a
    Free TPrimExp Int64 VName
    p) ExtLMAD
    lmad
    offsetMemoryInBodyReturns RebaseMap
    _ BranchTypeMem
    br = BranchTypeMem
    br

    offsetMemoryInLambda :: RebaseMap -> Lambda GPUMem -> OffsetM (Lambda GPUMem)
    offsetMemoryInLambda :: RebaseMap -> Lambda GPUMem -> OffsetM (Lambda GPUMem)
    offsetMemoryInLambda RebaseMap
    offsets Lambda GPUMem
    lam = do
      body <- Lambda GPUMem -> OffsetM (Body GPUMem) -> OffsetM (Body GPUMem)
    forall rep a (m :: * -> *) b.
    (Scoped rep a, LocalScope rep m) =>
    a -> m b -> m b
    inScopeOf Lambda GPUMem
    lam (OffsetM (Body GPUMem) -> OffsetM (Body GPUMem))
    -> OffsetM (Body GPUMem) -> OffsetM (Body GPUMem)
    forall a b. (a -> b) -> a -> b
    $ RebaseMap -> Body GPUMem -> OffsetM (Body GPUMem)
    offsetMemoryInBody RebaseMap
    offsets (Body GPUMem -> OffsetM (Body GPUMem))
    -> Body GPUMem -> OffsetM (Body GPUMem)
    forall a b. (a -> b) -> a -> b
    $ Lambda GPUMem -> Body GPUMem
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda GPUMem
    lam
      let params = (Param (MemInfo SubExp NoUniqueness MemBind)
     -> Param (MemInfo SubExp NoUniqueness MemBind))
    -> [Param (MemInfo SubExp NoUniqueness MemBind)]
    -> [Param (MemInfo SubExp NoUniqueness MemBind)]
    forall a b. (a -> b) -> [a] -> [b]
    map (RebaseMap
    -> Param (MemInfo SubExp NoUniqueness MemBind)
    -> Param (MemInfo SubExp NoUniqueness MemBind)
    forall u. RebaseMap -> Param (MemBound u) -> Param (MemBound u)
    offsetMemoryInParam RebaseMap
    offsets) ([Param (MemInfo SubExp NoUniqueness MemBind)]
     -> [Param (MemInfo SubExp NoUniqueness MemBind)])
    -> [Param (MemInfo SubExp NoUniqueness MemBind)]
    -> [Param (MemInfo SubExp NoUniqueness MemBind)]
    forall a b. (a -> b) -> a -> b
    $ Lambda GPUMem -> [LParam GPUMem]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda GPUMem
    lam
      pure $ lam {lambdaBody = body, lambdaParams = params}

    -- A loop may have memory parameters, and those memory blocks may
    -- be expanded.  We assume (but do not check - FIXME) that if the
    -- initial value of a loop parameter is an expanded memory block,
    -- then so will the result be.
    offsetMemoryInLoopParams ::
      RebaseMap ->
      [(FParam GPUMem, SubExp)] ->
      (RebaseMap -> [(FParam GPUMem, SubExp)] -> OffsetM a) ->
      OffsetM a
    offsetMemoryInLoopParams :: forall a.
    RebaseMap
    -> [(FParam GPUMem, SubExp)]
    -> (RebaseMap -> [(FParam GPUMem, SubExp)] -> OffsetM a)
    -> OffsetM a
    offsetMemoryInLoopParams RebaseMap
    offsets [(FParam GPUMem, SubExp)]
    merge RebaseMap -> [(FParam GPUMem, SubExp)] -> OffsetM a
    f = do
      let ([Param (MemInfo SubExp Uniqueness MemBind)]
    params, [SubExp]
    args) = [(Param (MemInfo SubExp Uniqueness MemBind), SubExp)]
    -> ([Param (MemInfo SubExp Uniqueness MemBind)], [SubExp])
    forall a b. [(a, b)] -> ([a], [b])
    unzip [(FParam GPUMem, SubExp)]
    [(Param (MemInfo SubExp Uniqueness MemBind), SubExp)]
    merge
      params' <- [Param (MemInfo SubExp Uniqueness MemBind)]
    -> OffsetM [Param (MemInfo SubExp Uniqueness MemBind)]
    addParamsContext [Param (MemInfo SubExp Uniqueness MemBind)]
    params
      args' <- (args <>) <$> argsContext args
      f offsets' $ zip params' args'
      where
        offsets' :: RebaseMap
    offsets' = RebaseMap -> RebaseMap
    extend RebaseMap
    offsets
        extend :: RebaseMap -> RebaseMap
    extend RebaseMap
    rm = (RebaseMap
     -> (Param (MemInfo SubExp Uniqueness MemBind), SubExp)
     -> RebaseMap)
    -> RebaseMap
    -> [(Param (MemInfo SubExp Uniqueness MemBind), SubExp)]
    -> RebaseMap
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' RebaseMap
    -> (Param (MemInfo SubExp Uniqueness MemBind), SubExp) -> RebaseMap
    forall {a} {dec}. Map VName a -> (Param dec, SubExp) -> Map VName a
    onParamArg RebaseMap
    rm [(FParam GPUMem, SubExp)]
    [(Param (MemInfo SubExp Uniqueness MemBind), SubExp)]
    merge
        onParamArg :: Map VName a -> (Param dec, SubExp) -> Map VName a
    onParamArg Map VName a
    rm (Param dec
    param, Var VName
    arg)
          | Just a
    x <- VName -> Map VName a -> Maybe a
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    arg Map VName a
    rm =
              VName -> a -> Map VName a -> Map VName a
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert (Param dec -> VName
    forall dec. Param dec -> VName
    paramName Param dec
    param) a
    x Map VName a
    rm
        onParamArg Map VName a
    rm (Param dec, SubExp)
    _ = Map VName a
    rm

    -- | Handles only the expressions where we do not change the number of
    -- results; meaning anything except Loop, Match, and nonscalar Apply.
    offsetMemoryInExp :: RebaseMap -> Exp GPUMem -> OffsetM (Exp GPUMem)
    offsetMemoryInExp :: RebaseMap -> Exp GPUMem -> OffsetM (Exp GPUMem)
    offsetMemoryInExp RebaseMap
    offsets = Mapper GPUMem GPUMem OffsetM -> Exp GPUMem -> OffsetM (Exp GPUMem)
    forall (m :: * -> *) frep trep.
    Monad m =>
    Mapper frep trep m -> Exp frep -> m (Exp trep)
    mapExpM Mapper GPUMem GPUMem OffsetM
    recurse
      where
        recurse :: Mapper GPUMem GPUMem OffsetM
    recurse =
          (forall rep (m :: * -> *). Monad m => Mapper rep rep m
    identityMapper @GPUMem)
            { mapOnBody = \Scope GPUMem
    bscope -> Scope GPUMem -> OffsetM (Body GPUMem) -> OffsetM (Body GPUMem)
    forall a. Scope GPUMem -> OffsetM a -> OffsetM a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope Scope GPUMem
    bscope (OffsetM (Body GPUMem) -> OffsetM (Body GPUMem))
    -> (Body GPUMem -> OffsetM (Body GPUMem))
    -> Body GPUMem
    -> OffsetM (Body GPUMem)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . RebaseMap -> Body GPUMem -> OffsetM (Body GPUMem)
    offsetMemoryInBody RebaseMap
    offsets,
              mapOnBranchType = pure . offsetMemoryInBodyReturns offsets,
              mapOnOp = onOp
            }
        onOp :: MemOp (HostOp NoOp) GPUMem -> OffsetM (MemOp (HostOp NoOp) GPUMem)
    onOp (Inner (SegOp SegOp SegLevel GPUMem
    op)) =
          HostOp NoOp GPUMem -> MemOp (HostOp NoOp) GPUMem
    forall (inner :: * -> *) rep. inner rep -> MemOp inner rep
    Inner (HostOp NoOp GPUMem -> MemOp (HostOp NoOp) GPUMem)
    -> (SegOp SegLevel GPUMem -> HostOp NoOp GPUMem)
    -> SegOp SegLevel GPUMem
    -> MemOp (HostOp NoOp) GPUMem
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegOp SegLevel GPUMem -> HostOp NoOp GPUMem
    forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
    SegOp
            (SegOp SegLevel GPUMem -> MemOp (HostOp NoOp) GPUMem)
    -> OffsetM (SegOp SegLevel GPUMem)
    -> OffsetM (MemOp (HostOp NoOp) GPUMem)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Scope GPUMem
    -> OffsetM (SegOp SegLevel GPUMem)
    -> OffsetM (SegOp SegLevel GPUMem)
    forall a. Scope GPUMem -> OffsetM a -> OffsetM a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope (SegSpace -> Scope GPUMem
    forall rep. SegSpace -> Scope rep
    scopeOfSegSpace (SegOp SegLevel GPUMem -> SegSpace
    forall lvl rep. SegOp lvl rep -> SegSpace
    segSpace SegOp SegLevel GPUMem
    op)) (SegOpMapper SegLevel GPUMem GPUMem OffsetM
    -> SegOp SegLevel GPUMem -> OffsetM (SegOp SegLevel GPUMem)
    forall (m :: * -> *) lvl frep trep.
    Monad m =>
    SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
    mapSegOpM SegOpMapper SegLevel GPUMem GPUMem OffsetM
    segOpMapper SegOp SegLevel GPUMem
    op)
          where
            segOpMapper :: SegOpMapper SegLevel GPUMem GPUMem OffsetM
    segOpMapper =
              SegOpMapper SegLevel (ZonkAny 1) (ZonkAny 1) OffsetM
    forall (m :: * -> *) lvl rep. Monad m => SegOpMapper lvl rep rep m
    identitySegOpMapper
                { mapOnSegOpBody = offsetMemoryInKernelBody offsets,
                  mapOnSegOpLambda = offsetMemoryInLambda offsets
                }
        onOp MemOp (HostOp NoOp) GPUMem
    op = MemOp (HostOp NoOp) GPUMem -> OffsetM (MemOp (HostOp NoOp) GPUMem)
    forall a. a -> OffsetM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure MemOp (HostOp NoOp) GPUMem
    op

    offsetMemoryInStm :: RebaseMap -> Stm GPUMem -> OffsetM (Stm GPUMem)
    offsetMemoryInStm :: RebaseMap -> Stm GPUMem -> OffsetM (Stm GPUMem)
    offsetMemoryInStm RebaseMap
    offsets (Let Pat (LetDec GPUMem)
    pat StmAux (ExpDec GPUMem)
    dec (Match [SubExp]
    cond [Case (Body GPUMem)]
    cases Body GPUMem
    defbody (MatchDec [BranchType GPUMem]
    ts MatchSort
    kind))) = do
      cases' <- [Case (Body GPUMem)]
    -> (Case (Body GPUMem) -> OffsetM (Case (Body GPUMem)))
    -> OffsetM [Case (Body GPUMem)]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [Case (Body GPUMem)]
    cases ((Case (Body GPUMem) -> OffsetM (Case (Body GPUMem)))
     -> OffsetM [Case (Body GPUMem)])
    -> (Case (Body GPUMem) -> OffsetM (Case (Body GPUMem)))
    -> OffsetM [Case (Body GPUMem)]
    forall a b. (a -> b) -> a -> b
    $ \(Case [Maybe PrimValue]
    vs Body GPUMem
    body) ->
        [Maybe PrimValue] -> Body GPUMem -> Case (Body GPUMem)
    forall body. [Maybe PrimValue] -> body -> Case body
    Case [Maybe PrimValue]
    vs (Body GPUMem -> Case (Body GPUMem))
    -> OffsetM (Body GPUMem) -> OffsetM (Case (Body GPUMem))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> RebaseMap -> Body GPUMem -> OffsetM (Body GPUMem)
    offsetMemoryInBodyReturnCtx RebaseMap
    offsets Body GPUMem
    body
      defbody' <- offsetMemoryInBodyReturnCtx offsets defbody
      (pat', ts') <- offsetBranch pat ts
      pure $ Let pat' dec $ Match cond cases' defbody' $ MatchDec ts' kind
    offsetMemoryInStm RebaseMap
    offsets (Let Pat (LetDec GPUMem)
    pat StmAux (ExpDec GPUMem)
    dec (Loop [(FParam GPUMem, SubExp)]
    merge LoopForm
    form Body GPUMem
    body)) = do
      loop' <-
        RebaseMap
    -> [(FParam GPUMem, SubExp)]
    -> (RebaseMap -> [(FParam GPUMem, SubExp)] -> OffsetM (Exp GPUMem))
    -> OffsetM (Exp GPUMem)
    forall a.
    RebaseMap
    -> [(FParam GPUMem, SubExp)]
    -> (RebaseMap -> [(FParam GPUMem, SubExp)] -> OffsetM a)
    -> OffsetM a
    offsetMemoryInLoopParams RebaseMap
    offsets [(FParam GPUMem, SubExp)]
    merge ((RebaseMap -> [(FParam GPUMem, SubExp)] -> OffsetM (Exp GPUMem))
     -> OffsetM (Exp GPUMem))
    -> (RebaseMap -> [(FParam GPUMem, SubExp)] -> OffsetM (Exp GPUMem))
    -> OffsetM (Exp GPUMem)
    forall a b. (a -> b) -> a -> b
    $ \RebaseMap
    offsets' [(FParam GPUMem, SubExp)]
    merge' -> do
          body' <-
            Scope GPUMem -> OffsetM (Body GPUMem) -> OffsetM (Body GPUMem)
    forall a. Scope GPUMem -> OffsetM a -> OffsetM a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope
              ([FParam GPUMem] -> Scope GPUMem
    forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfFParams (((FParam GPUMem, SubExp) -> FParam GPUMem)
    -> [(FParam GPUMem, SubExp)] -> [FParam GPUMem]
    forall a b. (a -> b) -> [a] -> [b]
    map (FParam GPUMem, SubExp) -> FParam GPUMem
    forall a b. (a, b) -> a
    fst [(FParam GPUMem, SubExp)]
    merge') Scope GPUMem -> Scope GPUMem -> Scope GPUMem
    forall a. Semigroup a => a -> a -> a
    <> LoopForm -> Scope GPUMem
    forall rep. LoopForm -> Scope rep
    scopeOfLoopForm LoopForm
    form)
              (RebaseMap -> Body GPUMem -> OffsetM (Body GPUMem)
    offsetMemoryInBodyReturnCtx RebaseMap
    offsets' Body GPUMem
    body)
          pure $ Loop merge' form body'
      pat' <- addPatternContext pat
      pure $ Let pat' dec loop'
    offsetMemoryInStm RebaseMap
    offsets (Let Pat (LetDec GPUMem)
    pat StmAux (ExpDec GPUMem)
    dec Exp GPUMem
    e) = do
      e' <- RebaseMap -> Exp GPUMem -> OffsetM (Exp GPUMem)
    offsetMemoryInExp RebaseMap
    offsets Exp GPUMem
    e
      pat' <-
        offsetMemoryInPat offsets pat
          <$> ( maybe (throwError "offsetMemoryInStm: ill-typed") pure
                  =<< expReturns e'
              )
      scope <- askScope
      -- Try to recompute the index function.  Fall back to creating rebase
      -- operations with the RebaseMap.
      rts <-
        maybe (throwError "offsetMemoryInStm: ill-typed") pure $
          runReader (expReturns e') scope
      let pat'' = [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    -> Pat (MemInfo SubExp NoUniqueness MemBind)
    forall dec. [PatElem dec] -> Pat dec
    Pat ([PatElem (MemInfo SubExp NoUniqueness MemBind)]
     -> Pat (MemInfo SubExp NoUniqueness MemBind))
    -> [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    -> Pat (MemInfo SubExp NoUniqueness MemBind)
    forall a b. (a -> b) -> a -> b
    $ (PatElem (MemInfo SubExp NoUniqueness MemBind)
     -> ExpReturns -> PatElem (MemInfo SubExp NoUniqueness MemBind))
    -> [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    -> [ExpReturns]
    -> [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith PatElem (MemInfo SubExp NoUniqueness MemBind)
    -> ExpReturns -> PatElem (MemInfo SubExp NoUniqueness MemBind)
    forall {d} {u} {d} {u}.
    PatElem (MemInfo d u MemBind)
    -> MemInfo d u (Maybe MemReturn) -> PatElem (MemInfo d u MemBind)
    pick (Pat (MemInfo SubExp NoUniqueness MemBind)
    -> [PatElem (MemInfo SubExp NoUniqueness MemBind)]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat (MemInfo SubExp NoUniqueness MemBind)
    pat') [ExpReturns]
    rts
      pure $ Let pat'' dec e'
      where
        pick :: PatElem (MemInfo d u MemBind)
    -> MemInfo d u (Maybe MemReturn) -> PatElem (MemInfo d u MemBind)
    pick
          (PatElem VName
    name (MemArray PrimType
    pt ShapeBase d
    s u
    u MemBind
    _ret))
          (MemArray PrimType
    _ ShapeBase d
    _ u
    _ (Just (ReturnsInBlock VName
    m ExtLMAD
    extlmad)))
            | Just LMAD
    lmad <- ExtLMAD -> Maybe LMAD
    instantiateLMAD ExtLMAD
    extlmad =
                VName -> MemInfo d u MemBind -> PatElem (MemInfo d u MemBind)
    forall dec. VName -> dec -> PatElem dec
    PatElem VName
    name (PrimType -> ShapeBase d -> u -> MemBind -> MemInfo d u MemBind
    forall d u ret.
    PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
    MemArray PrimType
    pt ShapeBase d
    s u
    u (VName -> LMAD -> MemBind
    ArrayIn VName
    m LMAD
    lmad))
        pick PatElem (MemInfo d u MemBind)
    p MemInfo d u (Maybe MemReturn)
    _ = PatElem (MemInfo d u MemBind)
    p

        instantiateLMAD :: ExtLMAD -> Maybe LMAD
        instantiateLMAD :: ExtLMAD -> Maybe LMAD
    instantiateLMAD = (TPrimExp Int64 (Ext VName) -> Maybe (TPrimExp Int64 VName))
    -> ExtLMAD -> Maybe LMAD
    forall (t :: * -> *) (f :: * -> *) a b.
    (Traversable t, Applicative f) =>
    (a -> f b) -> t a -> f (t b)
    forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> LMAD a -> f (LMAD b)
    traverse ((Ext VName -> Maybe VName)
    -> TPrimExp Int64 (Ext VName) -> Maybe (TPrimExp Int64 VName)
    forall (t :: * -> *) (f :: * -> *) a b.
    (Traversable t, Applicative f) =>
    (a -> f b) -> t a -> f (t b)
    forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> TPrimExp Int64 a -> f (TPrimExp Int64 b)
    traverse Ext VName -> Maybe VName
    forall {a}. Ext a -> Maybe a
    inst)
          where
            inst :: Ext a -> Maybe a
    inst Ext {} = Maybe a
    forall a. Maybe a
    Nothing
            inst (Free a
    x) = a -> Maybe a
    forall a. a -> Maybe a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure a
    x

    ---- Slicing allocation sizes out of a kernel.

    unAllocGPUStms :: Stms GPUMem -> Either String (Stms GPU.GPU)
    unAllocGPUStms :: Stms GPUMem -> Either String (Stms GPU)
    unAllocGPUStms = Bool -> Stms GPUMem -> Either String (Stms GPU)
    unAllocStms Bool
    False
      where
        unAllocBody :: Body GPUMem -> Either String (Body GPU)
    unAllocBody (Body BodyDec GPUMem
    dec Stms GPUMem
    stms Result
    res) =
          BodyDec GPU -> Stms GPU -> Result -> Body GPU
    forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
    Body BodyDec GPU
    BodyDec GPUMem
    dec (Stms GPU -> Result -> Body GPU)
    -> Either String (Stms GPU) -> Either String (Result -> Body GPU)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Bool -> Stms GPUMem -> Either String (Stms GPU)
    unAllocStms Bool
    True Stms GPUMem
    stms Either String (Result -> Body GPU)
    -> Either String Result -> Either String (Body GPU)
    forall a b.
    Either String (a -> b) -> Either String a -> Either String b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Result -> Either String Result
    forall a. a -> Either String a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Result
    res

        unAllocKernelBody :: KernelBody GPUMem -> Either String (KernelBody GPU)
    unAllocKernelBody (KernelBody BodyDec GPUMem
    dec Stms GPUMem
    stms [KernelResult]
    res) =
          BodyDec GPU -> Stms GPU -> [KernelResult] -> KernelBody GPU
    forall rep.
    BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
    KernelBody BodyDec GPU
    BodyDec GPUMem
    dec (Stms GPU -> [KernelResult] -> KernelBody GPU)
    -> Either String (Stms GPU)
    -> Either String ([KernelResult] -> KernelBody GPU)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Bool -> Stms GPUMem -> Either String (Stms GPU)
    unAllocStms Bool
    True Stms GPUMem
    stms Either String ([KernelResult] -> KernelBody GPU)
    -> Either String [KernelResult] -> Either String (KernelBody GPU)
    forall a b.
    Either String (a -> b) -> Either String a -> Either String b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> [KernelResult] -> Either String [KernelResult]
    forall a. a -> Either String a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [KernelResult]
    res

        unAllocStms :: Bool -> Stms GPUMem -> Either String (Stms GPU)
    unAllocStms Bool
    nested = (Stm GPUMem -> Either String (Stm GPU))
    -> Stms GPUMem -> Either String (Stms GPU)
    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) -> Seq a -> m (Seq b)
    mapM (Bool -> Stm GPUMem -> Either String (Stm GPU)
    unAllocStm Bool
    nested)

        unAllocStm :: Bool -> Stm GPUMem -> Either String (Stm GPU)
    unAllocStm Bool
    nested stm :: Stm GPUMem
    stm@(Let Pat (LetDec GPUMem)
    pat StmAux (ExpDec GPUMem)
    dec (Op Alloc {}))
          | Bool
    nested =
              String -> Either String (Stm GPU)
    forall a. String -> Either String a
    forall e (m :: * -> *) a. MonadError e m => e -> m a
    throwError (String -> Either String (Stm GPU))
    -> String -> Either String (Stm GPU)
    forall a b. (a -> b) -> a -> b
    $ String
    "Cannot handle nested allocation: " String -> String -> String
    forall a. Semigroup a => a -> a -> a
    <> Stm GPUMem -> String
    forall a. Pretty a => a -> String
    prettyString Stm GPUMem
    stm
          | Bool
    otherwise =
              Pat (TypeBase (ShapeBase SubExp) NoUniqueness)
    -> StmAux () -> Exp GPU -> Stm GPU
    Pat (LetDec GPU) -> StmAux (ExpDec GPU) -> Exp GPU -> Stm GPU
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let
                (Pat (TypeBase (ShapeBase SubExp) NoUniqueness)
     -> StmAux () -> Exp GPU -> Stm GPU)
    -> Either String (Pat (TypeBase (ShapeBase SubExp) NoUniqueness))
    -> Either String (StmAux () -> Exp GPU -> Stm GPU)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Pat (MemInfo SubExp NoUniqueness MemBind)
    -> Either String (Pat (TypeBase (ShapeBase SubExp) NoUniqueness))
    forall {d} {u} {ret} {a}.
    Pat (MemInfo d u ret) -> Either a (Pat (TypeBase (ShapeBase d) u))
    unAllocPat Pat (LetDec GPUMem)
    Pat (MemInfo SubExp NoUniqueness MemBind)
    pat
                Either String (StmAux () -> Exp GPU -> Stm GPU)
    -> Either String (StmAux ()) -> Either String (Exp GPU -> Stm GPU)
    forall a b.
    Either String (a -> b) -> Either String a -> Either String b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> StmAux () -> Either String (StmAux ())
    forall a. a -> Either String a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure StmAux ()
    StmAux (ExpDec GPUMem)
    dec
                Either String (Exp GPU -> Stm GPU)
    -> Either String (Exp GPU) -> Either String (Stm GPU)
    forall a b.
    Either String (a -> b) -> Either String a -> Either String b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Exp GPU -> Either String (Exp GPU)
    forall a. a -> Either String a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (BasicOp -> Exp GPU
    forall rep. BasicOp -> Exp rep
    BasicOp (SubExp -> BasicOp
    SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ PrimValue -> SubExp
    Constant PrimValue
    UnitValue))
        unAllocStm Bool
    _ (Let Pat (LetDec GPUMem)
    pat StmAux (ExpDec GPUMem)
    dec Exp GPUMem
    e) =
          Pat (TypeBase (ShapeBase SubExp) NoUniqueness)
    -> StmAux () -> Exp GPU -> Stm GPU
    Pat (LetDec GPU) -> StmAux (ExpDec GPU) -> Exp GPU -> Stm GPU
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let (Pat (TypeBase (ShapeBase SubExp) NoUniqueness)
     -> StmAux () -> Exp GPU -> Stm GPU)
    -> Either String (Pat (TypeBase (ShapeBase SubExp) NoUniqueness))
    -> Either String (StmAux () -> Exp GPU -> Stm GPU)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Pat (MemInfo SubExp NoUniqueness MemBind)
    -> Either String (Pat (TypeBase (ShapeBase SubExp) NoUniqueness))
    forall {d} {u} {ret} {a}.
    Pat (MemInfo d u ret) -> Either a (Pat (TypeBase (ShapeBase d) u))
    unAllocPat Pat (LetDec GPUMem)
    Pat (MemInfo SubExp NoUniqueness MemBind)
    pat Either String (StmAux () -> Exp GPU -> Stm GPU)
    -> Either String (StmAux ()) -> Either String (Exp GPU -> Stm GPU)
    forall a b.
    Either String (a -> b) -> Either String a -> Either String b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> StmAux () -> Either String (StmAux ())
    forall a. a -> Either String a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure StmAux ()
    StmAux (ExpDec GPUMem)
    dec Either String (Exp GPU -> Stm GPU)
    -> Either String (Exp GPU) -> Either String (Stm GPU)
    forall a b.
    Either String (a -> b) -> Either String a -> Either String b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Mapper GPUMem GPU (Either String)
    -> Exp GPUMem -> Either String (Exp GPU)
    forall (m :: * -> *) frep trep.
    Monad m =>
    Mapper frep trep m -> Exp frep -> m (Exp trep)
    mapExpM Mapper GPUMem GPU (Either String)
    unAlloc' Exp GPUMem
    e

        unAllocLambda :: Lambda GPUMem -> Either String (Lambda GPU)
    unAllocLambda (Lambda [LParam GPUMem]
    params [TypeBase (ShapeBase SubExp) NoUniqueness]
    ret Body GPUMem
    body) =
          [LParam GPU]
    -> [TypeBase (ShapeBase SubExp) NoUniqueness]
    -> Body GPU
    -> Lambda GPU
    forall rep.
    [LParam rep]
    -> [TypeBase (ShapeBase SubExp) NoUniqueness]
    -> Body rep
    -> Lambda rep
    Lambda ((Param (MemInfo SubExp NoUniqueness MemBind)
     -> Param (TypeBase (ShapeBase SubExp) NoUniqueness))
    -> [Param (MemInfo SubExp NoUniqueness MemBind)]
    -> [Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (MemInfo SubExp NoUniqueness MemBind)
    -> Param (TypeBase (ShapeBase SubExp) NoUniqueness)
    forall {d} {u} {ret}.
    Param (MemInfo d u ret) -> Param (TypeBase (ShapeBase d) u)
    unParam [LParam GPUMem]
    [Param (MemInfo SubExp NoUniqueness MemBind)]
    params) [TypeBase (ShapeBase SubExp) NoUniqueness]
    ret (Body GPU -> Lambda GPU)
    -> Either String (Body GPU) -> Either String (Lambda GPU)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Body GPUMem -> Either String (Body GPU)
    unAllocBody Body GPUMem
    body

        unAllocPat :: Pat (MemInfo d u ret) -> Either a (Pat (TypeBase (ShapeBase d) u))
    unAllocPat (Pat [PatElem (MemInfo d u ret)]
    pes) =
          [PatElem (TypeBase (ShapeBase d) u)]
    -> Pat (TypeBase (ShapeBase d) u)
    forall dec. [PatElem dec] -> Pat dec
    Pat ([PatElem (TypeBase (ShapeBase d) u)]
     -> Pat (TypeBase (ShapeBase d) u))
    -> Either a [PatElem (TypeBase (ShapeBase d) u)]
    -> Either a (Pat (TypeBase (ShapeBase d) u))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (PatElem (MemInfo d u ret)
     -> Either a (PatElem (TypeBase (ShapeBase d) u)))
    -> [PatElem (MemInfo d u ret)]
    -> Either a [PatElem (TypeBase (ShapeBase d) u)]
    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 ((MemInfo d u ret -> Either a (TypeBase (ShapeBase d) u))
    -> PatElem (MemInfo d u ret)
    -> Either a (PatElem (TypeBase (ShapeBase d) u))
    forall (m :: * -> *) from to.
    Monad m =>
    (from -> m to) -> PatElem from -> m (PatElem to)
    rephrasePatElem (TypeBase (ShapeBase d) u -> Either a (TypeBase (ShapeBase d) u)
    forall a b. b -> Either a b
    Right (TypeBase (ShapeBase d) u -> Either a (TypeBase (ShapeBase d) u))
    -> (MemInfo d u ret -> TypeBase (ShapeBase d) u)
    -> MemInfo d u ret
    -> Either a (TypeBase (ShapeBase d) u)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . MemInfo d u ret -> TypeBase (ShapeBase d) u
    forall d u ret. MemInfo d u ret -> TypeBase (ShapeBase d) u
    unMem)) [PatElem (MemInfo d u ret)]
    pes

        unAllocOp :: MemOp (HostOp NoOp) GPUMem -> Either String (HostOp SOAC GPU)
    unAllocOp Alloc {} = String -> Either String (HostOp SOAC GPU)
    forall a b. a -> Either a b
    Left String
    "unAllocOp: unhandled Alloc"
        unAllocOp (Inner OtherOp {}) = String -> Either String (HostOp SOAC GPU)
    forall a b. a -> Either a b
    Left String
    "unAllocOp: unhandled OtherOp"
        unAllocOp (Inner GPUBody {}) = String -> Either String (HostOp SOAC GPU)
    forall a b. a -> Either a b
    Left String
    "unAllocOp: unhandled GPUBody"
        unAllocOp (Inner (SizeOp SizeOp
    op)) = HostOp SOAC GPU -> Either String (HostOp SOAC GPU)
    forall a. a -> Either String a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (HostOp SOAC GPU -> Either String (HostOp SOAC GPU))
    -> HostOp SOAC GPU -> Either String (HostOp SOAC GPU)
    forall a b. (a -> b) -> a -> b
    $ SizeOp -> HostOp SOAC GPU
    forall (op :: * -> *) rep. SizeOp -> HostOp op rep
    SizeOp SizeOp
    op
        unAllocOp (Inner (SegOp SegOp SegLevel GPUMem
    op)) = SegOp SegLevel GPU -> HostOp SOAC GPU
    forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
    SegOp (SegOp SegLevel GPU -> HostOp SOAC GPU)
    -> Either String (SegOp SegLevel GPU)
    -> Either String (HostOp SOAC GPU)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SegOpMapper SegLevel GPUMem GPU (Either String)
    -> SegOp SegLevel GPUMem -> Either String (SegOp SegLevel GPU)
    forall (m :: * -> *) lvl frep trep.
    Monad m =>
    SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
    mapSegOpM SegOpMapper SegLevel GPUMem GPU (Either String)
    mapper SegOp SegLevel GPUMem
    op
          where
            mapper :: SegOpMapper SegLevel GPUMem GPU (Either String)
    mapper =
              SegOpMapper SegLevel (ZonkAny 0) (ZonkAny 0) (Either String)
    forall (m :: * -> *) lvl rep. Monad m => SegOpMapper lvl rep rep m
    identitySegOpMapper
                { mapOnSegOpLambda = unAllocLambda,
                  mapOnSegOpBody = unAllocKernelBody
                }

        unParam :: Param (MemInfo d u ret) -> Param (TypeBase (ShapeBase d) u)
    unParam = (MemInfo d u ret -> TypeBase (ShapeBase d) u)
    -> Param (MemInfo d u ret) -> Param (TypeBase (ShapeBase d) u)
    forall a b. (a -> b) -> Param a -> Param b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap MemInfo d u ret -> TypeBase (ShapeBase d) u
    forall d u ret. MemInfo d u ret -> TypeBase (ShapeBase d) u
    unMem

        unT :: MemInfo d u ret -> Either a (TypeBase (ShapeBase d) u)
    unT = TypeBase (ShapeBase d) u -> Either a (TypeBase (ShapeBase d) u)
    forall a b. b -> Either a b
    Right (TypeBase (ShapeBase d) u -> Either a (TypeBase (ShapeBase d) u))
    -> (MemInfo d u ret -> TypeBase (ShapeBase d) u)
    -> MemInfo d u ret
    -> Either a (TypeBase (ShapeBase d) u)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . MemInfo d u ret -> TypeBase (ShapeBase d) u
    forall d u ret. MemInfo d u ret -> TypeBase (ShapeBase d) u
    unMem

        unAlloc' :: Mapper GPUMem GPU (Either String)
    unAlloc' =
          Mapper
            { mapOnBody :: Scope GPU -> Body GPUMem -> Either String (Body GPU)
    mapOnBody = (Body GPUMem -> Either String (Body GPU))
    -> Scope GPU -> Body GPUMem -> Either String (Body GPU)
    forall a b. a -> b -> a
    const Body GPUMem -> Either String (Body GPU)
    unAllocBody,
              mapOnRetType :: RetType GPUMem -> Either String (RetType GPU)
    mapOnRetType = RetType GPUMem -> Either String (RetType GPU)
    RetTypeMem -> Either String DeclExtType
    forall {d} {u} {ret} {a}.
    MemInfo d u ret -> Either a (TypeBase (ShapeBase d) u)
    unT,
              mapOnBranchType :: BranchType GPUMem -> Either String (BranchType GPU)
    mapOnBranchType = BranchType GPUMem -> Either String (BranchType GPU)
    BranchTypeMem -> Either String ExtType
    forall {d} {u} {ret} {a}.
    MemInfo d u ret -> Either a (TypeBase (ShapeBase d) u)
    unT,
              mapOnFParam :: FParam GPUMem -> Either String (FParam GPU)
    mapOnFParam = Param (TypeBase (ShapeBase SubExp) Uniqueness)
    -> Either String (Param (TypeBase (ShapeBase SubExp) Uniqueness))
    forall a b. b -> Either a b
    Right (Param (TypeBase (ShapeBase SubExp) Uniqueness)
     -> Either String (Param (TypeBase (ShapeBase SubExp) Uniqueness)))
    -> (Param (MemInfo SubExp Uniqueness MemBind)
        -> Param (TypeBase (ShapeBase SubExp) Uniqueness))
    -> Param (MemInfo SubExp Uniqueness MemBind)
    -> Either String (Param (TypeBase (ShapeBase SubExp) Uniqueness))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Param (MemInfo SubExp Uniqueness MemBind)
    -> Param (TypeBase (ShapeBase SubExp) Uniqueness)
    forall {d} {u} {ret}.
    Param (MemInfo d u ret) -> Param (TypeBase (ShapeBase d) u)
    unParam,
              mapOnLParam :: LParam GPUMem -> Either String (LParam GPU)
    mapOnLParam = Param (TypeBase (ShapeBase SubExp) NoUniqueness)
    -> Either String (Param (TypeBase (ShapeBase SubExp) NoUniqueness))
    forall a b. b -> Either a b
    Right (Param (TypeBase (ShapeBase SubExp) NoUniqueness)
     -> Either
          String (Param (TypeBase (ShapeBase SubExp) NoUniqueness)))
    -> (Param (MemInfo SubExp NoUniqueness MemBind)
        -> Param (TypeBase (ShapeBase SubExp) NoUniqueness))
    -> Param (MemInfo SubExp NoUniqueness MemBind)
    -> Either String (Param (TypeBase (ShapeBase SubExp) NoUniqueness))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Param (MemInfo SubExp NoUniqueness MemBind)
    -> Param (TypeBase (ShapeBase SubExp) NoUniqueness)
    forall {d} {u} {ret}.
    Param (MemInfo d u ret) -> Param (TypeBase (ShapeBase d) u)
    unParam,
              mapOnOp :: Op GPUMem -> Either String (Op GPU)
    mapOnOp = Op GPUMem -> Either String (Op GPU)
    MemOp (HostOp NoOp) GPUMem -> Either String (HostOp SOAC GPU)
    unAllocOp,
              mapOnSubExp :: SubExp -> Either String SubExp
    mapOnSubExp = SubExp -> Either String SubExp
    forall a b. b -> Either a b
    Right,
              mapOnVName :: VName -> Either String VName
    mapOnVName = VName -> Either String VName
    forall a b. b -> Either a b
    Right
            }

    unMem :: MemInfo d u ret -> TypeBase (ShapeBase d) u
    unMem :: forall d u ret. MemInfo d u ret -> TypeBase (ShapeBase d) u
    unMem (MemPrim PrimType
    pt) = PrimType -> TypeBase (ShapeBase d) u
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    pt
    unMem (MemArray PrimType
    pt ShapeBase d
    shape u
    u ret
    _) = PrimType -> ShapeBase d -> u -> TypeBase (ShapeBase d) u
    forall shape u. PrimType -> shape -> u -> TypeBase shape u
    Array PrimType
    pt ShapeBase d
    shape u
    u
    unMem (MemAcc VName
    acc ShapeBase SubExp
    ispace [TypeBase (ShapeBase SubExp) NoUniqueness]
    ts u
    u) = VName
    -> ShapeBase SubExp
    -> [TypeBase (ShapeBase SubExp) NoUniqueness]
    -> u
    -> TypeBase (ShapeBase d) u
    forall shape u.
    VName
    -> ShapeBase SubExp
    -> [TypeBase (ShapeBase SubExp) NoUniqueness]
    -> u
    -> TypeBase shape u
    Acc VName
    acc ShapeBase SubExp
    ispace [TypeBase (ShapeBase SubExp) NoUniqueness]
    ts u
    u
    unMem MemMem {} = PrimType -> TypeBase (ShapeBase d) u
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    Unit

    unAllocScope :: Scope GPUMem -> Scope GPU.GPU
    unAllocScope :: Scope GPUMem -> Scope GPU
    unAllocScope = (NameInfo GPUMem -> NameInfo GPU) -> Scope GPUMem -> Scope GPU
    forall a b k. (a -> b) -> Map k a -> Map k b
    M.map NameInfo GPUMem -> NameInfo GPU
    forall {rep} {d} {u} {rep} {ret} {d} {u} {ret} {d} {u} {ret}.
    (LetDec rep ~ TypeBase (ShapeBase d) u,
     LetDec rep ~ MemInfo d u ret,
     FParamInfo rep ~ TypeBase (ShapeBase d) u,
     FParamInfo rep ~ MemInfo d u ret,
     LParamInfo rep ~ TypeBase (ShapeBase d) u,
     LParamInfo rep ~ MemInfo d u ret) =>
    NameInfo rep -> NameInfo rep
    unInfo
      where
        unInfo :: NameInfo rep -> NameInfo rep
    unInfo (LetName LetDec rep
    dec) = LetDec rep -> NameInfo rep
    forall rep. LetDec rep -> NameInfo rep
    LetName (LetDec rep -> NameInfo rep) -> LetDec rep -> NameInfo rep
    forall a b. (a -> b) -> a -> b
    $ MemInfo d u ret -> TypeBase (ShapeBase d) u
    forall d u ret. MemInfo d u ret -> TypeBase (ShapeBase d) u
    unMem LetDec rep
    MemInfo d u ret
    dec
        unInfo (FParamName FParamInfo rep
    dec) = FParamInfo rep -> NameInfo rep
    forall rep. FParamInfo rep -> NameInfo rep
    FParamName (FParamInfo rep -> NameInfo rep) -> FParamInfo rep -> NameInfo rep
    forall a b. (a -> b) -> a -> b
    $ MemInfo d u ret -> TypeBase (ShapeBase d) u
    forall d u ret. MemInfo d u ret -> TypeBase (ShapeBase d) u
    unMem FParamInfo rep
    MemInfo d u ret
    dec
        unInfo (LParamName LParamInfo rep
    dec) = LParamInfo rep -> NameInfo rep
    forall rep. LParamInfo rep -> NameInfo rep
    LParamName (LParamInfo rep -> NameInfo rep) -> LParamInfo rep -> NameInfo rep
    forall a b. (a -> b) -> a -> b
    $ MemInfo d u ret -> TypeBase (ShapeBase d) u
    forall d u ret. MemInfo d u ret -> TypeBase (ShapeBase d) u
    unMem LParamInfo rep
    MemInfo d u ret
    dec
        unInfo (IndexName IntType
    it) = IntType -> NameInfo rep
    forall rep. IntType -> NameInfo rep
    IndexName IntType
    it

    removeCommonSizes :: Extraction -> [(SubExp, [(VName, Space)])]
    removeCommonSizes :: Extraction -> [(SubExp, [(VName, Space)])]
    removeCommonSizes = Map SubExp [(VName, Space)] -> [(SubExp, [(VName, Space)])]
    forall k a. Map k a -> [(k, a)]
    M.toList (Map SubExp [(VName, Space)] -> [(SubExp, [(VName, Space)])])
    -> (Extraction -> Map SubExp [(VName, Space)])
    -> Extraction
    -> [(SubExp, [(VName, Space)])]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Map SubExp [(VName, Space)]
     -> (VName, ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space))
     -> Map SubExp [(VName, Space)])
    -> Map SubExp [(VName, Space)]
    -> [(VName, ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space))]
    -> Map SubExp [(VName, Space)]
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' Map SubExp [(VName, Space)]
    -> (VName, ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space))
    -> Map SubExp [(VName, Space)]
    forall {k} {a} {b} {a}.
    Ord k =>
    Map k [(a, b)] -> (a, (a, k, b)) -> Map k [(a, b)]
    comb Map SubExp [(VName, Space)]
    forall a. Monoid a => a
    mempty ([(VName, ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space))]
     -> Map SubExp [(VName, Space)])
    -> (Extraction
        -> [(VName, ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space))])
    -> Extraction
    -> Map SubExp [(VName, Space)]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Extraction
    -> [(VName, ((SegLevel, [TPrimExp Int64 VName]), SubExp, Space))]
    forall k a. Map k a -> [(k, a)]
    M.toList
      where
        comb :: Map k [(a, b)] -> (a, (a, k, b)) -> Map k [(a, b)]
    comb Map k [(a, b)]
    m (a
    mem, (a
    _, k
    size, b
    space)) = ([(a, b)] -> [(a, b)] -> [(a, b)])
    -> k -> [(a, b)] -> Map k [(a, b)] -> Map k [(a, b)]
    forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
    M.insertWith [(a, b)] -> [(a, b)] -> [(a, b)]
    forall a. [a] -> [a] -> [a]
    (++) k
    size [(a
    mem, b
    space)] Map k [(a, b)]
    m

    copyConsumed :: (MonadBuilder m, AliasableRep (Rep m)) => Stms (Rep m) -> m (Stms (Rep m))
    copyConsumed :: forall (m :: * -> *).
    (MonadBuilder m, AliasableRep (Rep m)) =>
    Stms (Rep m) -> m (Stms (Rep m))
    copyConsumed Stms (Rep m)
    stms = do
      let consumed :: [VName]
    consumed = Names -> [VName]
    namesToList (Names -> [VName]) -> Names -> [VName]
    forall a b. (a -> b) -> a -> b
    $ (AliasTable, Names) -> Names
    forall a b. (a, b) -> b
    snd ((AliasTable, Names) -> Names) -> (AliasTable, Names) -> Names
    forall a b. (a -> b) -> a -> b
    $ (Stms (Aliases (Rep m)), (AliasTable, Names))
    -> (AliasTable, Names)
    forall a b. (a, b) -> b
    snd ((Stms (Aliases (Rep m)), (AliasTable, Names))
     -> (AliasTable, Names))
    -> (Stms (Aliases (Rep m)), (AliasTable, Names))
    -> (AliasTable, Names)
    forall a b. (a -> b) -> a -> b
    $ AliasTable
    -> Stms (Rep m) -> (Stms (Aliases (Rep m)), (AliasTable, Names))
    forall rep.
    AliasableRep rep =>
    AliasTable -> Stms rep -> (Stms (Aliases rep), (AliasTable, Names))
    Alias.analyseStms AliasTable
    forall a. Monoid a => a
    mempty Stms (Rep m)
    stms
      m () -> m (Stms (Rep m))
    forall (m :: * -> *) a. MonadBuilder m => m a -> m (Stms (Rep m))
    collectStms_ (m () -> m (Stms (Rep m))) -> m () -> m (Stms (Rep m))
    forall a b. (a -> b) -> a -> b
    $ do
        consumed' <- (VName -> m VName) -> [VName] -> m [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 -> m VName
    forall {m :: * -> *}. MonadBuilder m => VName -> m VName
    copy [VName]
    consumed
        let substs = [(VName, VName)] -> Map VName VName
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList ([VName] -> [VName] -> [(VName, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    consumed [VName]
    consumed')
        addStms $ substituteNames substs stms
      where
        copy :: VName -> m VName
    copy VName
    v = String -> Exp (Rep m) -> m VName
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m VName
    letExp (VName -> String
    baseString VName
    v String -> String -> String
    forall a. Semigroup a => a -> a -> a
    <> String
    "_copy") (Exp (Rep m) -> m VName) -> Exp (Rep m) -> m VName
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
    forall a b. (a -> b) -> a -> b
    $ ShapeBase SubExp -> SubExp -> BasicOp
    Replicate ShapeBase SubExp
    forall a. Monoid a => a
    mempty (SubExp -> BasicOp) -> SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var VName
    v

    -- Important for edge cases (#1838) that the Stms here still have the
    -- Allocs we are actually trying to get rid of.
    sliceKernelSizes ::
      SubExp ->
      [SubExp] ->
      SegSpace ->
      Stms GPUMem ->
      ExpandM (Stms GPU.GPU, [VName], [VName])
    sliceKernelSizes :: SubExp
    -> [SubExp]
    -> SegSpace
    -> Stms GPUMem
    -> ExpandM (Stms GPU, [VName], [VName])
    sliceKernelSizes SubExp
    num_threads [SubExp]
    sizes SegSpace
    space Stms GPUMem
    kstms = do
      kstms' <- (String
     -> ReaderT
          (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPU))
    -> (Stms GPU
        -> ReaderT
             (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPU))
    -> Either String (Stms GPU)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPU)
    forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
    either String
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPU)
    forall a.
    String
    -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall e (m :: * -> *) a. MonadError e m => e -> m a
    throwError Stms GPU
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPU)
    forall a.
    a -> ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Either String (Stms GPU)
     -> ReaderT
          (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPU))
    -> Either String (Stms GPU)
    -> ReaderT
         (Scope GPUMem) (StateT VNameSource (Either String)) (Stms GPU)
    forall a b. (a -> b) -> a -> b
    $ Stms GPUMem -> Either String (Stms GPU)
    unAllocGPUStms Stms GPUMem
    kstms
      let num_sizes = [SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [SubExp]
    sizes
          i64s = Int
    -> TypeBase (ShapeBase SubExp) NoUniqueness
    -> [TypeBase (ShapeBase SubExp) NoUniqueness]
    forall a. Int -> a -> [a]
    replicate Int
    num_sizes (TypeBase (ShapeBase SubExp) NoUniqueness
     -> [TypeBase (ShapeBase SubExp) NoUniqueness])
    -> TypeBase (ShapeBase SubExp) NoUniqueness
    -> [TypeBase (ShapeBase SubExp) NoUniqueness]
    forall a b. (a -> b) -> a -> b
    $ PrimType -> TypeBase (ShapeBase SubExp) NoUniqueness
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64

      kernels_scope <- asks unAllocScope

      (max_lam, _) <- flip runBuilderT kernels_scope $ do
        xs <- replicateM num_sizes $ newParam "x" (Prim int64)
        ys <- replicateM num_sizes $ newParam "y" (Prim int64)
        (zs, stms) <- localScope (scopeOfLParams $ xs ++ ys) $
          collectStms $
            forM (zip xs ys) $ \(Param (TypeBase (ShapeBase SubExp) NoUniqueness)
    x, Param (TypeBase (ShapeBase SubExp) NoUniqueness)
    y) ->
              (SubExp -> SubExpRes)
    -> BuilderT
         GPU
         (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
         SubExp
    -> BuilderT
         GPU
         (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
         SubExpRes
    forall a b.
    (a -> b)
    -> BuilderT
         GPU (ReaderT (Scope GPUMem) (StateT VNameSource (Either String))) a
    -> BuilderT
         GPU (ReaderT (Scope GPUMem) (StateT VNameSource (Either String))) b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap SubExp -> SubExpRes
    subExpRes (BuilderT
       GPU
       (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
       SubExp
     -> BuilderT
          GPU
          (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
          SubExpRes)
    -> (BasicOp
        -> BuilderT
             GPU
             (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
             SubExp)
    -> BasicOp
    -> BuilderT
         GPU
         (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
         SubExpRes
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . String
    -> Exp
         (Rep
            (BuilderT
               GPU (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))))
    -> BuilderT
         GPU
         (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
         SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    "z" (Exp GPU
     -> BuilderT
          GPU
          (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
          SubExp)
    -> (BasicOp -> Exp GPU)
    -> BasicOp
    -> BuilderT
         GPU
         (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
         SubExp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . BasicOp -> Exp GPU
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp
     -> BuilderT
          GPU
          (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
          SubExpRes)
    -> BasicOp
    -> BuilderT
         GPU
         (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
         SubExpRes
    forall a b. (a -> b) -> a -> b
    $
                BinOp -> SubExp -> SubExp -> BasicOp
    BinOp (IntType -> BinOp
    SMax IntType
    Int64) (VName -> SubExp
    Var (VName -> SubExp) -> VName -> SubExp
    forall a b. (a -> b) -> a -> b
    $ Param (TypeBase (ShapeBase SubExp) NoUniqueness) -> VName
    forall dec. Param dec -> VName
    paramName Param (TypeBase (ShapeBase SubExp) NoUniqueness)
    x) (VName -> SubExp
    Var (VName -> SubExp) -> VName -> SubExp
    forall a b. (a -> b) -> a -> b
    $ Param (TypeBase (ShapeBase SubExp) NoUniqueness) -> VName
    forall dec. Param dec -> VName
    paramName Param (TypeBase (ShapeBase SubExp) NoUniqueness)
    y)
        pure $ Lambda (xs ++ ys) i64s (mkBody stms zs)

      flat_gtid_lparam <- newParam "flat_gtid" (Prim (IntType Int64))

      size_lam' <- localScope (scopeOfSegSpace space) . fmap fst . flip runBuilderT kernels_scope $
        GPU.simplifyLambda <=< mkLambda [flat_gtid_lparam] $ do
          -- Even though this SegRed is one-dimensional, we need to
          -- provide indexes corresponding to the original potentially
          -- multi-dimensional construct.
          let (kspace_gtids, kspace_dims) = unzip $ unSegSpace space
              new_inds =
                [TPrimExp Int64 VName]
    -> TPrimExp Int64 VName -> [TPrimExp Int64 VName]
    forall num. IntegralExp num => [num] -> num -> [num]
    unflattenIndex
                  ((SubExp -> TPrimExp Int64 VName)
    -> [SubExp] -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> TPrimExp Int64 VName
    pe64 [SubExp]
    kspace_dims)
                  (SubExp -> TPrimExp Int64 VName
    pe64 (SubExp -> TPrimExp Int64 VName) -> SubExp -> TPrimExp Int64 VName
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var (VName -> SubExp) -> VName -> SubExp
    forall a b. (a -> b) -> a -> b
    $ Param (TypeBase (ShapeBase SubExp) NoUniqueness) -> VName
    forall dec. Param dec -> VName
    paramName Param (TypeBase (ShapeBase SubExp) NoUniqueness)
    flat_gtid_lparam)
          zipWithM_ letBindNames (map pure kspace_gtids) =<< mapM toExp new_inds
          mapM_ addStm =<< copyConsumed kstms'
          pure $ subExpsRes sizes

      ((maxes_per_thread, size_sums), slice_stms) <- flip runBuilderT kernels_scope $ do
        pat <-
          basicPat <$> replicateM num_sizes (newIdent "max_per_thread" $ Prim int64)

        w <-
          letSubExp "size_slice_w"
            =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) (segSpaceDims space)

        thread_space_iota <-
          letExp "thread_space_iota" $
            BasicOp $
              Iota w (intConst Int64 0) (intConst Int64 1) Int64
        let red_op =
              Commutativity
    -> Lambda GPU -> [SubExp] -> ShapeBase SubExp -> SegBinOp GPU
    forall rep.
    Commutativity
    -> Lambda rep -> [SubExp] -> ShapeBase SubExp -> SegBinOp rep
    SegBinOp
                Commutativity
    Commutative
                Lambda GPU
    max_lam
                (Int -> SubExp -> [SubExp]
    forall a. Int -> a -> [a]
    replicate Int
    num_sizes (SubExp -> [SubExp]) -> SubExp -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    0)
                ShapeBase SubExp
    forall a. Monoid a => a
    mempty
        lvl <- segThread "segred"

        addStms
          =<< mapM renameStm
          =<< nonSegRed lvl pat w [red_op] size_lam' [thread_space_iota]

        size_sums <- forM (patNames pat) $ \VName
    threads_max ->
          String
    -> Exp
         (Rep
            (BuilderT
               GPU (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))))
    -> BuilderT
         GPU
         (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
         VName
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m VName
    letExp String
    "size_sum" (Exp
       (Rep
          (BuilderT
             GPU (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))))
     -> BuilderT
          GPU
          (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
          VName)
    -> Exp
         (Rep
            (BuilderT
               GPU (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))))
    -> BuilderT
         GPU
         (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))
         VName
    forall a b. (a -> b) -> a -> b
    $
            BasicOp
    -> Exp
         (Rep
            (BuilderT
               GPU (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))))
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp
     -> Exp
          (Rep
             (BuilderT
                GPU
                (ReaderT (Scope GPUMem) (StateT VNameSource (Either String))))))
    -> BasicOp
    -> Exp
         (Rep
            (BuilderT
               GPU (ReaderT (Scope GPUMem) (StateT VNameSource (Either String)))))
    forall a b. (a -> b) -> a -> b
    $
              BinOp -> SubExp -> SubExp -> BasicOp
    BinOp (IntType -> Overflow -> BinOp
    Mul IntType
    Int64 Overflow
    OverflowUndef) (VName -> SubExp
    Var VName
    threads_max) SubExp
    num_threads

        pure (patNames pat, size_sums)

      pure (slice_stms, maxes_per_thread, size_sums)