    {-# LANGUAGE TypeFamilies #-}

    -- | This pass attempts to lift allocations and asserts as far towards
    -- the top in their body as possible. This helps memory short
    -- circuiting do a better job, as it is sensitive to statement
    -- ordering.  It does not try to hoist allocations outside across body
    -- boundaries.
    module Futhark.Pass.LiftAllocations
      ( liftAllocationsSeqMem,
        liftAllocationsGPUMem,
        liftAllocationsMCMem,
      )
    where

    import Control.Monad.Reader
    import Data.Map qualified as M
    import Data.Sequence (Seq (..))
    import Futhark.Analysis.Alias (aliasAnalysis)
    import Futhark.IR.Aliases
    import Futhark.IR.GPUMem
    import Futhark.IR.MCMem
    import Futhark.IR.SeqMem
    import Futhark.Pass (Pass (..))

    liftInProg ::
      (AliasableRep rep, Mem rep inner, ASTConstraints (inner (Aliases rep))) =>
      (inner (Aliases rep) -> LiftM (inner (Aliases rep)) (inner (Aliases rep))) ->
      Prog rep ->
      Prog rep
    liftInProg :: forall rep (inner :: * -> *).
    (AliasableRep rep, Mem rep inner,
     ASTConstraints (inner (Aliases rep))) =>
    (inner (Aliases rep)
     -> LiftM (inner (Aliases rep)) (inner (Aliases rep)))
    -> Prog rep -> Prog rep
    liftInProg inner (Aliases rep)
    -> LiftM (inner (Aliases rep)) (inner (Aliases rep))
    onOp Prog rep
    prog =
      Prog rep
    prog
        { progFuns = removeFunDefAliases . onFun <$> progFuns (aliasAnalysis prog)
        }
      where
        onFun :: FunDef (Aliases rep) -> FunDef (Aliases rep)
    onFun FunDef (Aliases rep)
    f = FunDef (Aliases rep)
    f {funDefBody = onBody (funDefBody f)}
        onBody :: Body (Aliases rep) -> Body (Aliases rep)
    onBody Body (Aliases rep)
    body = Reader (Env (inner (Aliases rep))) (Body (Aliases rep))
    -> Env (inner (Aliases rep)) -> Body (Aliases rep)
    forall r a. Reader r a -> r -> a
    runReader (Body (Aliases rep)
    -> Reader (Env (inner (Aliases rep))) (Body (Aliases rep))
    forall rep (inner :: * -> *).
    (Mem rep inner, Aliased rep) =>
    Body rep -> LiftM (inner rep) (Body rep)
    liftAllocationsInBody Body (Aliases rep)
    body) ((inner (Aliases rep)
     -> LiftM (inner (Aliases rep)) (inner (Aliases rep)))
    -> AliasesAndConsumed -> Env (inner (Aliases rep))
    forall inner.
    (inner -> LiftM inner inner) -> AliasesAndConsumed -> Env inner
    Env inner (Aliases rep)
    -> LiftM (inner (Aliases rep)) (inner (Aliases rep))
    onOp AliasesAndConsumed
    forall a. Monoid a => a
    mempty)

    liftAllocationsSeqMem :: Pass SeqMem SeqMem
    liftAllocationsSeqMem :: Pass SeqMem SeqMem
    liftAllocationsSeqMem =
      String
    -> String
    -> (Prog SeqMem -> PassM (Prog SeqMem))
    -> Pass SeqMem SeqMem
    forall fromrep torep.
    String
    -> String
    -> (Prog fromrep -> PassM (Prog torep))
    -> Pass fromrep torep
    Pass String
    "lift allocations" String
    "lift allocations" ((Prog SeqMem -> PassM (Prog SeqMem)) -> Pass SeqMem SeqMem)
    -> (Prog SeqMem -> PassM (Prog SeqMem)) -> Pass SeqMem SeqMem
    forall a b. (a -> b) -> a -> b
    $
        Prog SeqMem -> PassM (Prog SeqMem)
    forall a. a -> PassM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Prog SeqMem -> PassM (Prog SeqMem))
    -> (Prog SeqMem -> Prog SeqMem)
    -> Prog SeqMem
    -> PassM (Prog SeqMem)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (NoOp (Aliases SeqMem)
     -> LiftM (NoOp (Aliases SeqMem)) (NoOp (Aliases SeqMem)))
    -> Prog SeqMem -> Prog SeqMem
    forall rep (inner :: * -> *).
    (AliasableRep rep, Mem rep inner,
     ASTConstraints (inner (Aliases rep))) =>
    (inner (Aliases rep)
     -> LiftM (inner (Aliases rep)) (inner (Aliases rep)))
    -> Prog rep -> Prog rep
    liftInProg NoOp (Aliases SeqMem)
    -> LiftM (NoOp (Aliases SeqMem)) (NoOp (Aliases SeqMem))
    forall a. a -> ReaderT (Env (NoOp (Aliases SeqMem))) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure

    liftAllocationsGPUMem :: Pass GPUMem GPUMem
    liftAllocationsGPUMem :: Pass GPUMem GPUMem
    liftAllocationsGPUMem =
      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
    "lift allocations gpu" String
    "lift allocations gpu" ((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 -> PassM (Prog GPUMem)
    forall a. a -> PassM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Prog GPUMem -> PassM (Prog GPUMem))
    -> (Prog GPUMem -> Prog GPUMem)
    -> Prog GPUMem
    -> PassM (Prog GPUMem)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (HostOp NoOp (Aliases GPUMem)
     -> LiftM
          (HostOp NoOp (Aliases GPUMem)) (HostOp NoOp (Aliases GPUMem)))
    -> Prog GPUMem -> Prog GPUMem
    forall rep (inner :: * -> *).
    (AliasableRep rep, Mem rep inner,
     ASTConstraints (inner (Aliases rep))) =>
    (inner (Aliases rep)
     -> LiftM (inner (Aliases rep)) (inner (Aliases rep)))
    -> Prog rep -> Prog rep
    liftInProg HostOp NoOp (Aliases GPUMem)
    -> LiftM
         (HostOp NoOp (Aliases GPUMem)) (HostOp NoOp (Aliases GPUMem))
    liftAllocationsInHostOp

    liftAllocationsMCMem :: Pass MCMem MCMem
    liftAllocationsMCMem :: Pass MCMem MCMem
    liftAllocationsMCMem =
      String
    -> String -> (Prog MCMem -> PassM (Prog MCMem)) -> Pass MCMem MCMem
    forall fromrep torep.
    String
    -> String
    -> (Prog fromrep -> PassM (Prog torep))
    -> Pass fromrep torep
    Pass String
    "lift allocations mc" String
    "lift allocations mc" ((Prog MCMem -> PassM (Prog MCMem)) -> Pass MCMem MCMem)
    -> (Prog MCMem -> PassM (Prog MCMem)) -> Pass MCMem MCMem
    forall a b. (a -> b) -> a -> b
    $
        Prog MCMem -> PassM (Prog MCMem)
    forall a. a -> PassM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Prog MCMem -> PassM (Prog MCMem))
    -> (Prog MCMem -> Prog MCMem) -> Prog MCMem -> PassM (Prog MCMem)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (MCOp NoOp (Aliases MCMem)
     -> LiftM (MCOp NoOp (Aliases MCMem)) (MCOp NoOp (Aliases MCMem)))
    -> Prog MCMem -> Prog MCMem
    forall rep (inner :: * -> *).
    (AliasableRep rep, Mem rep inner,
     ASTConstraints (inner (Aliases rep))) =>
    (inner (Aliases rep)
     -> LiftM (inner (Aliases rep)) (inner (Aliases rep)))
    -> Prog rep -> Prog rep
    liftInProg MCOp NoOp (Aliases MCMem)
    -> LiftM (MCOp NoOp (Aliases MCMem)) (MCOp NoOp (Aliases MCMem))
    liftAllocationsInMCOp

    data Env inner = Env
      { forall inner. Env inner -> inner -> LiftM inner inner
    onInner :: inner -> LiftM inner inner,
        forall inner. Env inner -> AliasesAndConsumed
    envAliases :: AliasesAndConsumed
      }

    type LiftM inner a = Reader (Env inner) a

    liftAllocationsInBody ::
      (Mem rep inner, Aliased rep) =>
      Body rep ->
      LiftM (inner rep) (Body rep)
    liftAllocationsInBody :: forall rep (inner :: * -> *).
    (Mem rep inner, Aliased rep) =>
    Body rep -> LiftM (inner rep) (Body rep)
    liftAllocationsInBody Body rep
    body = do
      stms <- Stms rep -> LiftM (inner rep) (Stms rep)
    forall rep (inner :: * -> *).
    (Mem rep inner, Aliased rep) =>
    Stms rep -> LiftM (inner rep) (Stms rep)
    liftAllocationsInStms (Body rep -> Stms rep
    forall rep. Body rep -> Stms rep
    bodyStms Body rep
    body)
      pure $ body {bodyStms = stms}

    liftInsideStm ::
      (Mem rep inner, Aliased rep) =>
      Stm rep ->
      LiftM (inner rep) (Stm rep)
    liftInsideStm :: forall rep (inner :: * -> *).
    (Mem rep inner, Aliased rep) =>
    Stm rep -> LiftM (inner rep) (Stm rep)
    liftInsideStm stm :: Stm rep
    stm@(Let Pat (LetDec rep)
    _ StmAux (ExpDec rep)
    _ (Op (Inner inner rep
    inner))) = do
      on_inner <- (Env (inner rep) -> inner rep -> LiftM (inner rep) (inner rep))
    -> ReaderT
         (Env (inner rep))
         Identity
         (inner rep -> LiftM (inner rep) (inner rep))
    forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
    asks Env (inner rep) -> inner rep -> LiftM (inner rep) (inner rep)
    forall inner. Env inner -> inner -> LiftM inner inner
    onInner
      inner' <- on_inner inner
      pure $ stm {stmExp = Op $ Inner inner'}
    liftInsideStm stm :: Stm rep
    stm@(Let Pat (LetDec rep)
    _ StmAux (ExpDec rep)
    _ (Match [SubExp]
    cond_ses [Case (Body rep)]
    cases Body rep
    body MatchDec (BranchType rep)
    dec)) = do
      cases' <- (Case (Body rep)
     -> ReaderT (Env (inner rep)) Identity (Case (Body rep)))
    -> [Case (Body rep)]
    -> ReaderT (Env (inner rep)) Identity [Case (Body rep)]
    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 [Maybe PrimValue]
    p Body rep
    b) -> [Maybe PrimValue] -> Body rep -> Case (Body rep)
    forall body. [Maybe PrimValue] -> body -> Case body
    Case [Maybe PrimValue]
    p (Body rep -> Case (Body rep))
    -> ReaderT (Env (inner rep)) Identity (Body rep)
    -> ReaderT (Env (inner rep)) Identity (Case (Body rep))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Body rep -> ReaderT (Env (inner rep)) Identity (Body rep)
    forall rep (inner :: * -> *).
    (Mem rep inner, Aliased rep) =>
    Body rep -> LiftM (inner rep) (Body rep)
    liftAllocationsInBody Body rep
    b) [Case (Body rep)]
    cases
      body' <- liftAllocationsInBody body
      pure stm {stmExp = Match cond_ses cases' body' dec}
    liftInsideStm stm :: Stm rep
    stm@(Let Pat (LetDec rep)
    _ StmAux (ExpDec rep)
    _ (Loop [(FParam rep, SubExp)]
    params LoopForm
    form Body rep
    body)) = do
      body' <- Body rep -> ReaderT (Env (inner rep)) Identity (Body rep)
    forall rep (inner :: * -> *).
    (Mem rep inner, Aliased rep) =>
    Body rep -> LiftM (inner rep) (Body rep)
    liftAllocationsInBody Body rep
    body
      pure stm {stmExp = Loop params form body'}
    liftInsideStm Stm rep
    stm = Stm rep -> ReaderT (Env (inner rep)) Identity (Stm rep)
    forall a. a -> ReaderT (Env (inner rep)) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Stm rep
    stm

    liftAllocationsInStms ::
      forall rep inner.
      (Mem rep inner, Aliased rep) =>
      Stms rep ->
      LiftM (inner rep) (Stms rep)
    liftAllocationsInStms :: forall rep (inner :: * -> *).
    (Mem rep inner, Aliased rep) =>
    Stms rep -> LiftM (inner rep) (Stms rep)
    liftAllocationsInStms Stms rep
    stms_orig = do
      outer_aliases <- (Env (inner rep) -> AliasesAndConsumed)
    -> ReaderT (Env (inner rep)) Identity AliasesAndConsumed
    forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
    asks Env (inner rep) -> AliasesAndConsumed
    forall inner. Env inner -> AliasesAndConsumed
    envAliases
      let aliases = (AliasesAndConsumed -> Stm rep -> AliasesAndConsumed)
    -> AliasesAndConsumed -> Stms rep -> AliasesAndConsumed
    forall b a. (b -> a -> b) -> b -> Seq a -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl AliasesAndConsumed -> Stm rep -> AliasesAndConsumed
    forall rep.
    Aliased rep =>
    AliasesAndConsumed -> Stm rep -> AliasesAndConsumed
    trackAliases AliasesAndConsumed
    outer_aliases Stms rep
    stms_orig
          go ::
            -- The input stms
            Stms rep ->
            -- The lifted allocations and associated statements
            Stms rep ->
            -- The other statements processed so far
            Stms rep ->
            -- (Names we need to lift, consumed names)
            (Names, Names) ->
            LiftM (inner rep) (Stms rep)
          go Stms rep
    Empty Stms rep
    lifted Stms rep
    acc (Names, Names)
    _ = Stms rep -> LiftM (inner rep) (Stms rep)
    forall a. a -> ReaderT (Env (inner rep)) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Stms rep -> LiftM (inner rep) (Stms rep))
    -> Stms rep -> LiftM (inner rep) (Stms rep)
    forall a b. (a -> b) -> a -> b
    $ Stms rep
    lifted Stms rep -> Stms rep -> Stms rep
    forall a. Semigroup a => a -> a -> a
    <> Stms rep
    acc
          go (Stms rep
    stms :|> Stm rep
    stm) Stms rep
    lifted Stms rep
    acc (Names
    to_lift, Names
    consumed) = do
            stm' <- Stm rep -> LiftM (inner rep) (Stm rep)
    forall rep (inner :: * -> *).
    (Mem rep inner, Aliased rep) =>
    Stm rep -> LiftM (inner rep) (Stm rep)
    liftInsideStm Stm rep
    stm
            case stmExp stm' of
              BasicOp Assert {} -> Stm rep -> LiftM (inner rep) (Stms rep)
    liftStm Stm rep
    stm'
              Op Alloc {} -> Stm rep -> LiftM (inner rep) (Stms rep)
    liftStm Stm rep
    stm'
              Exp rep
    _ -> do
                let pat_names :: Names
    pat_names = [VName] -> Names
    namesFromList ([VName] -> Names) -> [VName] -> Names
    forall a b. (a -> b) -> a -> b
    $ Pat (LetDec rep) -> [VName]
    forall dec. Pat dec -> [VName]
    patNames (Pat (LetDec rep) -> [VName]) -> Pat (LetDec rep) -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Stm rep -> Pat (LetDec rep)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm rep
    stm'
                    free_in_stm :: Names
    free_in_stm = Stm rep -> Names
    forall a. FreeIn a => a -> Names
    freeIn Stm rep
    stm
                    expand :: VName -> [VName]
    expand VName
    v = [VName] -> (Names -> [VName]) -> Maybe Names -> [VName]
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe [VName
    v] Names -> [VName]
    namesToList (Maybe Names -> [VName]) -> Maybe Names -> [VName]
    forall a b. (a -> b) -> a -> b
    $ VName -> Map VName Names -> Maybe Names
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    v (Map VName Names -> Maybe Names) -> Map VName Names -> Maybe Names
    forall a b. (a -> b) -> a -> b
    $ AliasesAndConsumed -> Map VName Names
    forall a b. (a, b) -> a
    fst AliasesAndConsumed
    aliases
                if (Names
    pat_names Names -> Names -> Bool
    `namesIntersect` Names
    to_lift)
                  Bool -> Bool -> Bool
    || (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (VName -> Names -> Bool
    `nameIn` Names
    free_in_stm) ((VName -> [VName]) -> [VName] -> [VName]
    forall m a. Monoid m => (a -> m) -> [a] -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap VName -> [VName]
    expand ([VName] -> [VName]) -> [VName] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Names -> [VName]
    namesToList Names
    consumed)
                  then Stm rep -> LiftM (inner rep) (Stms rep)
    liftStm Stm rep
    stm'
                  else Stm rep -> LiftM (inner rep) (Stms rep)
    dontLiftStm Stm rep
    stm'
            where
              liftStm :: Stm rep -> LiftM (inner rep) (Stms rep)
    liftStm Stm rep
    stm' =
                Stms rep
    -> Stms rep
    -> Stms rep
    -> (Names, Names)
    -> LiftM (inner rep) (Stms rep)
    go Stms rep
    stms (Stm rep
    stm' Stm rep -> Stms rep -> Stms rep
    forall a. a -> Seq a -> Seq a
    :<| Stms rep
    lifted) Stms rep
    acc (Names
    to_lift', Names
    consumed')
                where
                  to_lift' :: Names
    to_lift' =
                    Stm rep -> Names
    forall a. FreeIn a => a -> Names
    freeIn Stm rep
    stm'
                      Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> (Names
    to_lift Names -> Names -> Names
    `namesSubtract` [VName] -> Names
    namesFromList (Pat (LetDec rep) -> [VName]
    forall dec. Pat dec -> [VName]
    patNames (Stm rep -> Pat (LetDec rep)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm rep
    stm')))
                  consumed' :: Names
    consumed' = Names
    consumed Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Stm rep -> Names
    forall rep. Aliased rep => Stm rep -> Names
    consumedInStm Stm rep
    stm'
              dontLiftStm :: Stm rep -> LiftM (inner rep) (Stms rep)
    dontLiftStm Stm rep
    stm' =
                Stms rep
    -> Stms rep
    -> Stms rep
    -> (Names, Names)
    -> LiftM (inner rep) (Stms rep)
    go Stms rep
    stms Stms rep
    lifted (Stm rep
    stm' Stm rep -> Stms rep -> Stms rep
    forall a. a -> Seq a -> Seq a
    :<| Stms rep
    acc) (Names
    to_lift, Names
    consumed)
      local (\Env (inner rep)
    env -> Env (inner rep)
    env {envAliases = aliases}) $
        go stms_orig mempty mempty mempty

    liftAllocationsInSegOp ::
      (Mem rep inner, Aliased rep) =>
      SegOp lvl rep ->
      LiftM (inner rep) (SegOp lvl rep)
    liftAllocationsInSegOp :: forall rep (inner :: * -> *) lvl.
    (Mem rep inner, Aliased rep) =>
    SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep)
    liftAllocationsInSegOp (SegMap lvl
    lvl SegSpace
    sp [Type]
    tps KernelBody rep
    body) = do
      stms <- Stms rep -> LiftM (inner rep) (Stms rep)
    forall rep (inner :: * -> *).
    (Mem rep inner, Aliased rep) =>
    Stms rep -> LiftM (inner rep) (Stms rep)
    liftAllocationsInStms (KernelBody rep -> Stms rep
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody rep
    body)
      pure $ SegMap lvl sp tps $ body {kernelBodyStms = stms}
    liftAllocationsInSegOp (SegRed lvl
    lvl SegSpace
    sp [Type]
    tps KernelBody rep
    body [SegBinOp rep]
    binops) = do
      stms <- Stms rep -> LiftM (inner rep) (Stms rep)
    forall rep (inner :: * -> *).
    (Mem rep inner, Aliased rep) =>
    Stms rep -> LiftM (inner rep) (Stms rep)
    liftAllocationsInStms (KernelBody rep -> Stms rep
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody rep
    body)
      pure $ SegRed lvl sp tps (body {kernelBodyStms = stms}) binops
    liftAllocationsInSegOp (SegScan lvl
    lvl SegSpace
    sp [Type]
    tps KernelBody rep
    body [SegBinOp rep]
    binops) = do
      stms <- Stms rep -> LiftM (inner rep) (Stms rep)
    forall rep (inner :: * -> *).
    (Mem rep inner, Aliased rep) =>
    Stms rep -> LiftM (inner rep) (Stms rep)
    liftAllocationsInStms (KernelBody rep -> Stms rep
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody rep
    body)
      pure $ SegScan lvl sp tps (body {kernelBodyStms = stms}) binops
    liftAllocationsInSegOp (SegHist lvl
    lvl SegSpace
    sp [Type]
    tps KernelBody rep
    body [HistOp rep]
    histops) = do
      stms <- Stms rep -> LiftM (inner rep) (Stms rep)
    forall rep (inner :: * -> *).
    (Mem rep inner, Aliased rep) =>
    Stms rep -> LiftM (inner rep) (Stms rep)
    liftAllocationsInStms (KernelBody rep -> Stms rep
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody rep
    body)
      pure $ SegHist lvl sp tps (body {kernelBodyStms = stms}) histops

    liftAllocationsInHostOp ::
      HostOp NoOp (Aliases GPUMem) ->
      LiftM (HostOp NoOp (Aliases GPUMem)) (HostOp NoOp (Aliases GPUMem))
    liftAllocationsInHostOp :: HostOp NoOp (Aliases GPUMem)
    -> LiftM
         (HostOp NoOp (Aliases GPUMem)) (HostOp NoOp (Aliases GPUMem))
    liftAllocationsInHostOp (SegOp SegOp SegLevel (Aliases GPUMem)
    op) = SegOp SegLevel (Aliases GPUMem) -> HostOp NoOp (Aliases GPUMem)
    forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
    SegOp (SegOp SegLevel (Aliases GPUMem) -> HostOp NoOp (Aliases GPUMem))
    -> ReaderT
         (Env (HostOp NoOp (Aliases GPUMem)))
         Identity
         (SegOp SegLevel (Aliases GPUMem))
    -> LiftM
         (HostOp NoOp (Aliases GPUMem)) (HostOp NoOp (Aliases GPUMem))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SegOp SegLevel (Aliases GPUMem)
    -> ReaderT
         (Env (HostOp NoOp (Aliases GPUMem)))
         Identity
         (SegOp SegLevel (Aliases GPUMem))
    forall rep (inner :: * -> *) lvl.
    (Mem rep inner, Aliased rep) =>
    SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep)
    liftAllocationsInSegOp SegOp SegLevel (Aliases GPUMem)
    op
    liftAllocationsInHostOp HostOp NoOp (Aliases GPUMem)
    op = HostOp NoOp (Aliases GPUMem)
    -> LiftM
         (HostOp NoOp (Aliases GPUMem)) (HostOp NoOp (Aliases GPUMem))
    forall a.
    a -> ReaderT (Env (HostOp NoOp (Aliases GPUMem))) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure HostOp NoOp (Aliases GPUMem)
    op

    liftAllocationsInMCOp ::
      MCOp NoOp (Aliases MCMem) ->
      LiftM (MCOp NoOp (Aliases MCMem)) (MCOp NoOp (Aliases MCMem))
    liftAllocationsInMCOp :: MCOp NoOp (Aliases MCMem)
    -> LiftM (MCOp NoOp (Aliases MCMem)) (MCOp NoOp (Aliases MCMem))
    liftAllocationsInMCOp (ParOp Maybe (SegOp () (Aliases MCMem))
    par SegOp () (Aliases MCMem)
    op) =
      Maybe (SegOp () (Aliases MCMem))
    -> SegOp () (Aliases MCMem) -> MCOp NoOp (Aliases MCMem)
    forall (op :: * -> *) rep.
    Maybe (SegOp () rep) -> SegOp () rep -> MCOp op rep
    ParOp (Maybe (SegOp () (Aliases MCMem))
     -> SegOp () (Aliases MCMem) -> MCOp NoOp (Aliases MCMem))
    -> ReaderT
         (Env (MCOp NoOp (Aliases MCMem)))
         Identity
         (Maybe (SegOp () (Aliases MCMem)))
    -> ReaderT
         (Env (MCOp NoOp (Aliases MCMem)))
         Identity
         (SegOp () (Aliases MCMem) -> MCOp NoOp (Aliases MCMem))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (SegOp () (Aliases MCMem)
     -> ReaderT
          (Env (MCOp NoOp (Aliases MCMem)))
          Identity
          (SegOp () (Aliases MCMem)))
    -> Maybe (SegOp () (Aliases MCMem))
    -> ReaderT
         (Env (MCOp NoOp (Aliases MCMem)))
         Identity
         (Maybe (SegOp () (Aliases MCMem)))
    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) -> Maybe a -> f (Maybe b)
    traverse SegOp () (Aliases MCMem)
    -> ReaderT
         (Env (MCOp NoOp (Aliases MCMem)))
         Identity
         (SegOp () (Aliases MCMem))
    forall rep (inner :: * -> *) lvl.
    (Mem rep inner, Aliased rep) =>
    SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep)
    liftAllocationsInSegOp Maybe (SegOp () (Aliases MCMem))
    par ReaderT
      (Env (MCOp NoOp (Aliases MCMem)))
      Identity
      (SegOp () (Aliases MCMem) -> MCOp NoOp (Aliases MCMem))
    -> ReaderT
         (Env (MCOp NoOp (Aliases MCMem)))
         Identity
         (SegOp () (Aliases MCMem))
    -> LiftM (MCOp NoOp (Aliases MCMem)) (MCOp NoOp (Aliases MCMem))
    forall a b.
    ReaderT (Env (MCOp NoOp (Aliases MCMem))) Identity (a -> b)
    -> ReaderT (Env (MCOp NoOp (Aliases MCMem))) Identity a
    -> ReaderT (Env (MCOp NoOp (Aliases MCMem))) Identity b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SegOp () (Aliases MCMem)
    -> ReaderT
         (Env (MCOp NoOp (Aliases MCMem)))
         Identity
         (SegOp () (Aliases MCMem))
    forall rep (inner :: * -> *) lvl.
    (Mem rep inner, Aliased rep) =>
    SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep)
    liftAllocationsInSegOp SegOp () (Aliases MCMem)
    op
    liftAllocationsInMCOp MCOp NoOp (Aliases MCMem)
    op = MCOp NoOp (Aliases MCMem)
    -> LiftM (MCOp NoOp (Aliases MCMem)) (MCOp NoOp (Aliases MCMem))
    forall a. a -> ReaderT (Env (MCOp NoOp (Aliases MCMem))) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure MCOp NoOp (Aliases MCMem)
    op