{-# 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.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)))
-> Env (inner (Aliases rep))
forall inner. (inner -> LiftM inner inner) -> Env inner
Env inner (Aliases rep)
-> LiftM (inner (Aliases rep)) (inner (Aliases rep))
onOp)

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

newtype Env inner = Env
  {forall inner. Env inner -> inner -> LiftM inner inner
onInner :: inner -> LiftM inner inner}

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 rep
stms <- Stms rep
-> Stms rep
-> Stms rep
-> (Names, Names)
-> LiftM (inner rep) (Stms rep)
forall rep (inner :: * -> *).
(Mem rep inner, Aliased rep) =>
Stms rep
-> Stms rep
-> Stms rep
-> (Names, Names)
-> LiftM (inner rep) (Stms rep)
liftAllocationsInStms (Body rep -> Stms rep
forall rep. Body rep -> Stms rep
bodyStms Body rep
body) Stms rep
forall a. Monoid a => a
mempty Stms rep
forall a. Monoid a => a
mempty (Names, Names)
forall a. Monoid a => a
mempty
  Body rep -> LiftM (inner rep) (Body rep)
forall a. a -> ReaderT (Env (inner rep)) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Body rep -> LiftM (inner rep) (Body rep))
-> Body rep -> LiftM (inner rep) (Body rep)
forall a b. (a -> b) -> a -> b
$ Body rep
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
  inner rep -> LiftM (inner rep) (inner rep)
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 rep
inner' <- inner rep -> LiftM (inner rep) (inner rep)
on_inner inner rep
inner
  Stm rep -> LiftM (inner rep) (Stm rep)
forall a. a -> ReaderT (Env (inner rep)) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stm rep -> LiftM (inner rep) (Stm rep))
-> Stm rep -> LiftM (inner rep) (Stm rep)
forall a b. (a -> b) -> a -> b
$ Stm rep
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
  [Case (Body rep)]
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 rep
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
  Stm rep -> LiftM (inner rep) (Stm rep)
forall a. a -> ReaderT (Env (inner rep)) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Stm rep
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 rep
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
  Stm rep -> LiftM (inner rep) (Stm rep)
forall a. a -> ReaderT (Env (inner rep)) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Stm rep
stm {stmExp = Loop params form body'}
liftInsideStm Stm rep
stm = Stm rep -> LiftM (inner rep) (Stm rep)
forall a. a -> ReaderT (Env (inner rep)) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Stm rep
stm

liftAllocationsInStms ::
  (Mem rep inner, Aliased rep) =>
  -- | 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)
liftAllocationsInStms :: forall rep (inner :: * -> *).
(Mem rep inner, Aliased rep) =>
Stms rep
-> Stms rep
-> Stms rep
-> (Names, Names)
-> LiftM (inner rep) (Stms rep)
liftAllocationsInStms Stms rep
Empty Stms rep
lifted Stms rep
acc (Names, Names)
_ = Stms rep -> ReaderT (Env (inner rep)) Identity (Stms rep)
forall a. a -> ReaderT (Env (inner rep)) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms rep -> ReaderT (Env (inner rep)) Identity (Stms rep))
-> Stms rep -> ReaderT (Env (inner rep)) Identity (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
liftAllocationsInStms (Stms rep
stms :|> Stm rep
stm) Stms rep
lifted Stms rep
acc (Names
to_lift, Names
consumed) = do
  Stm rep
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 Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm' of
    BasicOp Assert {} -> Stm rep -> ReaderT (Env (inner rep)) Identity (Stms rep)
liftStm Stm rep
stm'
    Op Alloc {} -> Stm rep -> ReaderT (Env (inner rep)) Identity (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'
      if (Names
pat_names Names -> Names -> Bool
`namesIntersect` Names
to_lift)
        Bool -> Bool -> Bool
|| Names -> Names -> Bool
namesIntersect Names
consumed (Stm rep -> Names
forall a. FreeIn a => a -> Names
freeIn Stm rep
stm)
        then Stm rep -> ReaderT (Env (inner rep)) Identity (Stms rep)
liftStm Stm rep
stm'
        else Stm rep -> ReaderT (Env (inner rep)) Identity (Stms rep)
dontLiftStm Stm rep
stm'
  where
    liftStm :: Stm rep -> ReaderT (Env (inner rep)) Identity (Stms rep)
liftStm Stm rep
stm' =
      Stms rep
-> Stms rep
-> Stms rep
-> (Names, Names)
-> ReaderT (Env (inner rep)) Identity (Stms rep)
forall rep (inner :: * -> *).
(Mem rep inner, Aliased rep) =>
Stms rep
-> Stms rep
-> Stms rep
-> (Names, Names)
-> LiftM (inner rep) (Stms rep)
liftAllocationsInStms 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 -> ReaderT (Env (inner rep)) Identity (Stms rep)
dontLiftStm Stm rep
stm' =
      Stms rep
-> Stms rep
-> Stms rep
-> (Names, Names)
-> ReaderT (Env (inner rep)) Identity (Stms rep)
forall rep (inner :: * -> *).
(Mem rep inner, Aliased rep) =>
Stms rep
-> Stms rep
-> Stms rep
-> (Names, Names)
-> LiftM (inner rep) (Stms rep)
liftAllocationsInStms 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)

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 rep
stms <- Stms rep
-> Stms rep
-> Stms rep
-> (Names, Names)
-> LiftM (inner rep) (Stms rep)
forall rep (inner :: * -> *).
(Mem rep inner, Aliased rep) =>
Stms rep
-> Stms rep
-> Stms rep
-> (Names, Names)
-> LiftM (inner rep) (Stms rep)
liftAllocationsInStms (KernelBody rep -> Stms rep
forall rep. KernelBody rep -> Stms rep
kernelBodyStms KernelBody rep
body) Stms rep
forall a. Monoid a => a
mempty Stms rep
forall a. Monoid a => a
mempty (Names, Names)
forall a. Monoid a => a
mempty
  SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep)
forall a. a -> ReaderT (Env (inner rep)) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep))
-> SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep)
forall a b. (a -> b) -> a -> b
$ lvl -> SegSpace -> [Type] -> KernelBody rep -> SegOp lvl rep
forall lvl rep.
lvl -> SegSpace -> [Type] -> KernelBody rep -> SegOp lvl rep
SegMap lvl
lvl SegSpace
sp [Type]
tps (KernelBody rep -> SegOp lvl rep)
-> KernelBody rep -> SegOp lvl rep
forall a b. (a -> b) -> a -> b
$ KernelBody rep
body {kernelBodyStms = stms}
liftAllocationsInSegOp (SegRed lvl
lvl SegSpace
sp [Type]
tps KernelBody rep
body [SegBinOp rep]
binops) = do
  Stms rep
stms <- Stms rep
-> Stms rep
-> Stms rep
-> (Names, Names)
-> LiftM (inner rep) (Stms rep)
forall rep (inner :: * -> *).
(Mem rep inner, Aliased rep) =>
Stms rep
-> Stms rep
-> Stms rep
-> (Names, Names)
-> LiftM (inner rep) (Stms rep)
liftAllocationsInStms (KernelBody rep -> Stms rep
forall rep. KernelBody rep -> Stms rep
kernelBodyStms KernelBody rep
body) Stms rep
forall a. Monoid a => a
mempty Stms rep
forall a. Monoid a => a
mempty (Names, Names)
forall a. Monoid a => a
mempty
  SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep)
forall a. a -> ReaderT (Env (inner rep)) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep))
-> SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep)
forall a b. (a -> b) -> a -> b
$ lvl
-> SegSpace
-> [Type]
-> KernelBody rep
-> [SegBinOp rep]
-> SegOp lvl rep
forall lvl rep.
lvl
-> SegSpace
-> [Type]
-> KernelBody rep
-> [SegBinOp rep]
-> SegOp lvl rep
SegRed lvl
lvl SegSpace
sp [Type]
tps (KernelBody rep
body {kernelBodyStms = stms}) [SegBinOp rep]
binops
liftAllocationsInSegOp (SegScan lvl
lvl SegSpace
sp [Type]
tps KernelBody rep
body [SegBinOp rep]
binops) = do
  Stms rep
stms <- Stms rep
-> Stms rep
-> Stms rep
-> (Names, Names)
-> LiftM (inner rep) (Stms rep)
forall rep (inner :: * -> *).
(Mem rep inner, Aliased rep) =>
Stms rep
-> Stms rep
-> Stms rep
-> (Names, Names)
-> LiftM (inner rep) (Stms rep)
liftAllocationsInStms (KernelBody rep -> Stms rep
forall rep. KernelBody rep -> Stms rep
kernelBodyStms KernelBody rep
body) Stms rep
forall a. Monoid a => a
mempty Stms rep
forall a. Monoid a => a
mempty (Names, Names)
forall a. Monoid a => a
mempty
  SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep)
forall a. a -> ReaderT (Env (inner rep)) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep))
-> SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep)
forall a b. (a -> b) -> a -> b
$ lvl
-> SegSpace
-> [Type]
-> KernelBody rep
-> [SegBinOp rep]
-> SegOp lvl rep
forall lvl rep.
lvl
-> SegSpace
-> [Type]
-> KernelBody rep
-> [SegBinOp rep]
-> SegOp lvl rep
SegScan lvl
lvl SegSpace
sp [Type]
tps (KernelBody rep
body {kernelBodyStms = stms}) [SegBinOp rep]
binops
liftAllocationsInSegOp (SegHist lvl
lvl SegSpace
sp [Type]
tps KernelBody rep
body [HistOp rep]
histops) = do
  Stms rep
stms <- Stms rep
-> Stms rep
-> Stms rep
-> (Names, Names)
-> LiftM (inner rep) (Stms rep)
forall rep (inner :: * -> *).
(Mem rep inner, Aliased rep) =>
Stms rep
-> Stms rep
-> Stms rep
-> (Names, Names)
-> LiftM (inner rep) (Stms rep)
liftAllocationsInStms (KernelBody rep -> Stms rep
forall rep. KernelBody rep -> Stms rep
kernelBodyStms KernelBody rep
body) Stms rep
forall a. Monoid a => a
mempty Stms rep
forall a. Monoid a => a
mempty (Names, Names)
forall a. Monoid a => a
mempty
  SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep)
forall a. a -> ReaderT (Env (inner rep)) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep))
-> SegOp lvl rep -> LiftM (inner rep) (SegOp lvl rep)
forall a b. (a -> b) -> a -> b
$ lvl
-> SegSpace
-> [Type]
-> KernelBody rep
-> [HistOp rep]
-> SegOp lvl rep
forall lvl rep.
lvl
-> SegSpace
-> [Type]
-> KernelBody rep
-> [HistOp rep]
-> SegOp lvl rep
SegHist lvl
lvl SegSpace
sp [Type]
tps (KernelBody rep
body {kernelBodyStms = stms}) [HistOp rep]
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