{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE TypeFamilies #-}
module Futhark.Optimise.ArrayShortCircuiting.LastUse (lastUseSeqMem, lastUsePrg, lastUsePrgGPU, lastUseGPUMem) where
import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.Bifunctor (bimap)
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Sequence (Seq (..))
import Futhark.IR.Aliases
import Futhark.IR.GPUMem
import Futhark.IR.SeqMem
import Futhark.Optimise.ArrayShortCircuiting.DataStructs
import Futhark.Util
newtype LastUseReader rep = LastUseReader
{ forall {k} (rep :: k).
LastUseReader rep
-> Op (Aliases rep)
-> Names
-> LastUseM rep (AliasTab, Names, Names)
onOp :: Op (Aliases rep) -> Names -> LastUseM rep (LUTabFun, Names, Names)
}
type LastUseM rep a = StateT AliasTab (Reader (LastUseReader rep)) a
aliasLookup :: VName -> LastUseM rep Names
aliasLookup :: forall {k} (rep :: k). VName -> LastUseM rep Names
aliasLookup VName
vname =
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a -> a
fromMaybe forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
vname
lastUsePrg :: Prog (Aliases SeqMem) -> LUTabPrg
lastUsePrg :: Prog (Aliases SeqMem) -> LUTabPrg
lastUsePrg Prog (Aliases SeqMem)
prg = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map FunDef (Aliases SeqMem) -> (Name, AliasTab)
lastUseSeqMem forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Prog rep -> [FunDef rep]
progFuns Prog (Aliases SeqMem)
prg
lastUsePrgGPU :: Prog (Aliases GPUMem) -> LUTabPrg
lastUsePrgGPU :: Prog (Aliases GPUMem) -> LUTabPrg
lastUsePrgGPU Prog (Aliases GPUMem)
prg = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map FunDef (Aliases GPUMem) -> (Name, AliasTab)
lastUseGPUMem forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Prog rep -> [FunDef rep]
progFuns Prog (Aliases GPUMem)
prg
lastUseSeqMem :: FunDef (Aliases SeqMem) -> (Name, LUTabFun)
lastUseSeqMem :: FunDef (Aliases SeqMem) -> (Name, AliasTab)
lastUseSeqMem (FunDef Maybe EntryPoint
_ Attrs
_ Name
fname [RetType (Aliases SeqMem)]
_ [FParam (Aliases SeqMem)]
_ Body (Aliases SeqMem)
body) =
let (AliasTab
res, Names
_) =
forall r a. Reader r a -> r -> a
runReader
(forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseBody Body (Aliases SeqMem)
body (forall a. Monoid a => a
mempty, forall a. Monoid a => a
mempty)) forall a. Monoid a => a
mempty)
(forall {k} (rep :: k).
(Op (Aliases rep)
-> Names -> LastUseM rep (AliasTab, Names, Names))
-> LastUseReader rep
LastUseReader Op (Aliases SeqMem)
-> Names -> LastUseM SeqMem (AliasTab, Names, Names)
lastUseSeqOp)
in (Name
fname, AliasTab
res)
lastUseGPUMem :: FunDef (Aliases GPUMem) -> (Name, LUTabFun)
lastUseGPUMem :: FunDef (Aliases GPUMem) -> (Name, AliasTab)
lastUseGPUMem (FunDef Maybe EntryPoint
_ Attrs
_ Name
fname [RetType (Aliases GPUMem)]
_ [FParam (Aliases GPUMem)]
_ Body (Aliases GPUMem)
body) =
let (AliasTab
res, Names
_) =
forall r a. Reader r a -> r -> a
runReader
(forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseBody Body (Aliases GPUMem)
body (forall a. Monoid a => a
mempty, forall a. Monoid a => a
mempty)) forall a. Monoid a => a
mempty)
(forall {k} (rep :: k).
(Op (Aliases rep)
-> Names -> LastUseM rep (AliasTab, Names, Names))
-> LastUseReader rep
LastUseReader Op (Aliases GPUMem)
-> Names -> LastUseM GPUMem (AliasTab, Names, Names)
lastUseGPUOp)
in (Name
fname, AliasTab
res)
lastUseBody ::
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep) ->
(LUTabFun, Names) ->
LastUseM rep (LUTabFun, Names)
lastUseBody :: forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseBody bdy :: Body (Aliases rep)
bdy@(Body BodyDec (Aliases rep)
_ Stms (Aliases rep)
stms Result
result) (AliasTab
lutab, Names
used_nms) = do
(AliasTab
lutab', Names
_) <-
forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Stms (Aliases rep)
-> (AliasTab, Names) -> [VName] -> LastUseM rep (AliasTab, Names)
lastUseStms Stms (Aliases rep)
stms (AliasTab
lutab, Names
used_nms) forall a b. (a -> b) -> a -> b
$
Names -> [VName]
namesToList forall a b. (a -> b) -> a -> b
$
forall a. FreeIn a => a -> Names
freeIn forall a b. (a -> b) -> a -> b
$
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp Result
result
Names
used_in_body <- forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn Body (Aliases rep)
bdy
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
lutab', Names
used_nms forall a. Semigroup a => a -> a -> a
<> Names
used_in_body)
lastUseKernelBody ::
(CanBeAliased (Op rep), ASTRep rep) =>
KernelBody (Aliases rep) ->
(LUTabFun, Names) ->
LastUseM rep (LUTabFun, Names)
lastUseKernelBody :: forall {k} (rep :: k).
(CanBeAliased (Op rep), ASTRep rep) =>
KernelBody (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseKernelBody bdy :: KernelBody (Aliases rep)
bdy@(KernelBody BodyDec (Aliases rep)
_ Stms (Aliases rep)
stms [KernelResult]
result) (AliasTab
lutab, Names
used_nms) = do
(AliasTab
lutab', Names
_) <-
forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Stms (Aliases rep)
-> (AliasTab, Names) -> [VName] -> LastUseM rep (AliasTab, Names)
lastUseStms Stms (Aliases rep)
stms (AliasTab
lutab, Names
used_nms) forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn [KernelResult]
result
Names
used_in_body <- forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn KernelBody (Aliases rep)
bdy
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
lutab', Names
used_nms forall a. Semigroup a => a -> a -> a
<> Names
used_in_body)
lastUseStms ::
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Stms (Aliases rep) ->
(LUTabFun, Names) ->
[VName] ->
LastUseM rep (LUTabFun, Names)
lastUseStms :: forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Stms (Aliases rep)
-> (AliasTab, Names) -> [VName] -> LastUseM rep (AliasTab, Names)
lastUseStms Seq (Stm (Aliases rep))
Empty (AliasTab
lutab, Names
nms) [VName]
res_nms = do
Names
aliases <- forall (m :: * -> *) b a.
(Monad m, Monoid b) =>
(a -> m b) -> [a] -> m b
concatMapM forall {k} (rep :: k). VName -> LastUseM rep Names
aliasLookup [VName]
res_nms
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
lutab, Names
nms forall a. Semigroup a => a -> a -> a
<> Names
aliases)
lastUseStms (stm :: Stm (Aliases rep)
stm@(Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
_ Exp (Aliases rep)
e) :<| Seq (Stm (Aliases rep))
stms) (AliasTab
lutab, Names
nms) [VName]
res_nms = do
let extra_alias :: Names
extra_alias = case Exp (Aliases rep)
e of
BasicOp (Update Safety
_ VName
old Slice SubExp
_ SubExp
_) -> VName -> Names
oneName VName
old
BasicOp (FlatUpdate VName
old FlatSlice SubExp
_ VName
_) -> VName -> Names
oneName VName
old
Exp (Aliases rep)
_ -> forall a. Monoid a => a
mempty
forall {k} dec (rep :: k).
AliasesOf dec =>
Names -> Pat dec -> LastUseM rep ()
updateAliasing Names
extra_alias Pat (LetDec (Aliases rep))
pat
(AliasTab
lutab', Names
nms') <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Stms (Aliases rep)
-> (AliasTab, Names) -> [VName] -> LastUseM rep (AliasTab, Names)
lastUseStms Seq (Stm (Aliases rep))
stms (AliasTab
lutab, Names
nms) [VName]
res_nms
(AliasTab
lutab'', Names
nms'') <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Stm (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseStm Stm (Aliases rep)
stm (AliasTab
lutab', Names
nms')
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
lutab'', Names
nms'')
lastUseStm ::
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Stm (Aliases rep) ->
(LUTabFun, Names) ->
LastUseM rep (LUTabFun, Names)
lastUseStm :: forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Stm (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseStm (Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
_ Exp (Aliases rep)
e) (AliasTab
lutab, Names
used_nms) =
do
(AliasTab
lutab', Names
last_uses, Names
used_nms') <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Exp (Aliases rep) -> Names -> LastUseM rep (AliasTab, Names, Names)
lastUseExp Exp (Aliases rep)
e Names
used_nms
let patnms :: [VName]
patnms = forall dec. Pat dec -> [VName]
patNames Pat (LetDec (Aliases rep))
pat
used_nms'' :: Names
used_nms'' = Names
used_nms' Names -> Names -> Names
`namesSubtract` [VName] -> Names
namesFromList [VName]
patnms
lutab'' :: AliasTab
lutab'' =
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union AliasTab
lutab' forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (forall a. [a] -> a
head [VName]
patnms) Names
last_uses AliasTab
lutab
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
lutab'', Names
used_nms'')
lastUseExp ::
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Exp (Aliases rep) ->
Names ->
LastUseM rep (LUTabFun, Names, Names)
lastUseExp :: forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Exp (Aliases rep) -> Names -> LastUseM rep (AliasTab, Names, Names)
lastUseExp (Match [SubExp]
_ [Case (Body (Aliases rep))]
cases Body (Aliases rep)
body MatchDec (BranchType (Aliases rep))
_) Names
used_nms = do
(AliasTab
lutab_cases, Names
used_cases) <-
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap forall a. Monoid a => [a] -> a
mconcat forall a. Monoid a => [a] -> a
mconcat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [(a, b)] -> ([a], [b])
unzip
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseBody (forall k a. Map k a
M.empty, Names
used_nms) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall body. Case body -> body
caseBody) [Case (Body (Aliases rep))]
cases
(AliasTab
lutab', Names
body_used_nms) <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseBody Body (Aliases rep)
body (forall k a. Map k a
M.empty, Names
used_nms)
let free_in_body :: Names
free_in_body = forall a. FreeIn a => a -> Names
freeIn Body (Aliases rep)
body
let free_in_cases :: Names
free_in_cases = forall a. FreeIn a => a -> Names
freeIn [Case (Body (Aliases rep))]
cases
let used_nms' :: Names
used_nms' = Names
used_cases forall a. Semigroup a => a -> a -> a
<> Names
body_used_nms
(Names
_, Names
last_used_arrs) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms forall a b. (a -> b) -> a -> b
$ Names
free_in_body forall a. Semigroup a => a -> a -> a
<> Names
free_in_cases
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
lutab_cases forall a. Semigroup a => a -> a -> a
<> AliasTab
lutab', Names
last_used_arrs, Names
used_nms')
lastUseExp (DoLoop [(FParam (Aliases rep), SubExp)]
var_ses LoopForm (Aliases rep)
_ Body (Aliases rep)
body) Names
used_nms0 = do
Names
free_in_body <- forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn Body (Aliases rep)
body
[(VName, Names)]
var_inis <- forall a. [Maybe a] -> [a]
catMaybes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {dec} {rep :: k}.
Typed dec =>
Names
-> (Param dec, SubExp)
-> StateT
AliasTab (Reader (LastUseReader rep)) (Maybe (VName, Names))
initHelper (Names
free_in_body forall a. Semigroup a => a -> a -> a
<> Names
used_nms0)) [(FParam (Aliases rep), SubExp)]
var_ses
let
free_in_body' :: Names
free_in_body' = Names
free_in_body Names -> Names -> Names
`namesSubtract` [VName] -> Names
namesFromList (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(VName, Names)]
var_inis)
used_nms :: Names
used_nms = Names
used_nms0 forall a. Semigroup a => a -> a -> a
<> Names
free_in_body' forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn (forall {k} (rep :: k). Body rep -> Result
bodyResult Body (Aliases rep)
body)
(AliasTab
body_lutab, Names
_) <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseBody Body (Aliases rep)
body (forall a. Monoid a => a
mempty, Names
used_nms)
let lutab_res :: AliasTab
lutab_res = AliasTab
body_lutab forall a. Semigroup a => a -> a -> a
<> forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName, Names)]
var_inis
fpar_nms :: Names
fpar_nms = [VName] -> Names
namesFromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (Ident -> VName
identName forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. Typed dec => Param dec -> Ident
paramIdent forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(FParam (Aliases rep), SubExp)]
var_ses
used_nms' :: Names
used_nms' = (Names
free_in_body forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(FParam (Aliases rep), SubExp)]
var_ses)) Names -> Names -> Names
`namesSubtract` Names
fpar_nms
used_nms_res :: Names
used_nms_res = Names
used_nms0 forall a. Semigroup a => a -> a -> a
<> Names
used_nms' forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn (forall {k} (rep :: k). Body rep -> Result
bodyResult Body (Aliases rep)
body)
lu_arrs :: Names
lu_arrs = Names
used_nms' Names -> Names -> Names
`namesSubtract` Names
used_nms0
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
lutab_res, Names
lu_arrs, Names
used_nms_res)
where
initHelper :: Names
-> (Param dec, SubExp)
-> StateT
AliasTab (Reader (LastUseReader rep)) (Maybe (VName, Names))
initHelper Names
free_and_used (Param dec
fp, SubExp
se) = do
Names
names <- forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure forall a b. (a -> b) -> a -> b
$ forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty VName -> Names
oneName forall a b. (a -> b) -> a -> b
$ SubExp -> Maybe VName
subExpVar SubExp
se
if Names
names Names -> Names -> Bool
`namesIntersect` Names
free_and_used
then forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
else forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (Ident -> VName
identName forall a b. (a -> b) -> a -> b
$ forall dec. Typed dec => Param dec -> Ident
paramIdent Param dec
fp, Names
names)
lastUseExp (Op Op (Aliases rep)
op) Names
used_nms = do
OpWithAliases (Op rep)
-> Names -> LastUseM rep (AliasTab, Names, Names)
on_op <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
reader forall {k} (rep :: k).
LastUseReader rep
-> Op (Aliases rep)
-> Names
-> LastUseM rep (AliasTab, Names, Names)
onOp
OpWithAliases (Op rep)
-> Names -> LastUseM rep (AliasTab, Names, Names)
on_op Op (Aliases rep)
op Names
used_nms
lastUseExp Exp (Aliases rep)
e Names
used_nms = do
let free_in_e :: Names
free_in_e = forall a. FreeIn a => a -> Names
freeIn Exp (Aliases rep)
e
(Names
used_nms', Names
lu_vars) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms Names
free_in_e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall k a. Map k a
M.empty, Names
lu_vars, Names
used_nms')
lastUseGPUOp :: Op (Aliases GPUMem) -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)
lastUseGPUOp :: Op (Aliases GPUMem)
-> Names -> LastUseM GPUMem (AliasTab, Names, Names)
lastUseGPUOp (Alloc SubExp
se Space
sp) Names
used_nms = do
let free_in_e :: Names
free_in_e = forall a. FreeIn a => a -> Names
freeIn SubExp
se forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn Space
sp
(Names
used_nms', Names
lu_vars) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms Names
free_in_e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall k a. Map k a
M.empty, Names
lu_vars, Names
used_nms')
lastUseGPUOp (Inner (OtherOp ())) Names
used_nms =
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Monoid a => a
mempty, forall a. Monoid a => a
mempty, Names
used_nms)
lastUseGPUOp (Inner (SizeOp SizeOp
sop)) Names
used_nms = do
(Names
used_nms', Names
lu_vars) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn SizeOp
sop
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Monoid a => a
mempty, Names
lu_vars, Names
used_nms')
lastUseGPUOp (Inner (SegOp (SegMap SegLevel
_ SegSpace
_ [Type]
tps KernelBody (Aliases GPUMem)
kbody))) Names
used_nms = do
(Names
used_nms', Names
lu_vars) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn [Type]
tps
(AliasTab
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
(CanBeAliased (Op rep), ASTRep rep) =>
KernelBody (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseKernelBody KernelBody (Aliases GPUMem)
kbody (forall a. Monoid a => a
mempty, Names
used_nms')
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
body_lutab, Names
lu_vars, Names
used_nms' forall a. Semigroup a => a -> a -> a
<> Names
used_nms'')
lastUseGPUOp (Inner (SegOp (SegRed SegLevel
_ SegSpace
_ [SegBinOp (Aliases GPUMem)]
sbos [Type]
tps KernelBody (Aliases GPUMem)
kbody))) Names
used_nms = do
(AliasTab
lutab_sbo, Names
lu_vars_sbo, Names
used_nms_sbo) <- [SegBinOp (Aliases GPUMem)]
-> Names -> LastUseM GPUMem (AliasTab, Names, Names)
lastUseSegBinOp [SegBinOp (Aliases GPUMem)]
sbos Names
used_nms
(Names
used_nms', Names
lu_vars) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms_sbo forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn [Type]
tps
(AliasTab
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
(CanBeAliased (Op rep), ASTRep rep) =>
KernelBody (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseKernelBody KernelBody (Aliases GPUMem)
kbody (forall a. Monoid a => a
mempty, Names
used_nms')
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union AliasTab
lutab_sbo AliasTab
body_lutab, Names
lu_vars forall a. Semigroup a => a -> a -> a
<> Names
lu_vars_sbo, Names
used_nms_sbo forall a. Semigroup a => a -> a -> a
<> Names
used_nms' forall a. Semigroup a => a -> a -> a
<> Names
used_nms'')
lastUseGPUOp (Inner (SegOp (SegScan SegLevel
_ SegSpace
_ [SegBinOp (Aliases GPUMem)]
sbos [Type]
tps KernelBody (Aliases GPUMem)
kbody))) Names
used_nms = do
(AliasTab
lutab_sbo, Names
lu_vars_sbo, Names
used_nms_sbo) <- [SegBinOp (Aliases GPUMem)]
-> Names -> LastUseM GPUMem (AliasTab, Names, Names)
lastUseSegBinOp [SegBinOp (Aliases GPUMem)]
sbos Names
used_nms
(Names
used_nms', Names
lu_vars) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms_sbo forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn [Type]
tps
(AliasTab
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
(CanBeAliased (Op rep), ASTRep rep) =>
KernelBody (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseKernelBody KernelBody (Aliases GPUMem)
kbody (forall a. Monoid a => a
mempty, Names
used_nms')
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union AliasTab
lutab_sbo AliasTab
body_lutab, Names
lu_vars forall a. Semigroup a => a -> a -> a
<> Names
lu_vars_sbo, Names
used_nms_sbo forall a. Semigroup a => a -> a -> a
<> Names
used_nms' forall a. Semigroup a => a -> a -> a
<> Names
used_nms'')
lastUseGPUOp (Inner (SegOp (SegHist SegLevel
_ SegSpace
_ [HistOp (Aliases GPUMem)]
hos [Type]
tps KernelBody (Aliases GPUMem)
kbody))) Names
used_nms = do
(AliasTab
lutab_sbo, Names
lu_vars_sbo, Names
used_nms_sbo) <- [HistOp (Aliases GPUMem)]
-> Names -> LastUseM GPUMem (AliasTab, Names, Names)
lastUseHistOp [HistOp (Aliases GPUMem)]
hos Names
used_nms
(Names
used_nms', Names
lu_vars) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms_sbo forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn [Type]
tps
(AliasTab
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
(CanBeAliased (Op rep), ASTRep rep) =>
KernelBody (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseKernelBody KernelBody (Aliases GPUMem)
kbody (forall a. Monoid a => a
mempty, Names
used_nms')
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union AliasTab
lutab_sbo AliasTab
body_lutab, Names
lu_vars forall a. Semigroup a => a -> a -> a
<> Names
lu_vars_sbo, Names
used_nms_sbo forall a. Semigroup a => a -> a -> a
<> Names
used_nms' forall a. Semigroup a => a -> a -> a
<> Names
used_nms'')
lastUseGPUOp (Inner (GPUBody [Type]
tps Body (Aliases GPUMem)
body)) Names
used_nms = do
(Names
used_nms', Names
lu_vars) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn [Type]
tps
(AliasTab
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseBody Body (Aliases GPUMem)
body (forall a. Monoid a => a
mempty, Names
used_nms')
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
body_lutab, Names
lu_vars, Names
used_nms' forall a. Semigroup a => a -> a -> a
<> Names
used_nms'')
lastUseSegBinOp :: [SegBinOp (Aliases GPUMem)] -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)
lastUseSegBinOp :: [SegBinOp (Aliases GPUMem)]
-> Names -> LastUseM GPUMem (AliasTab, Names, Names)
lastUseSegBinOp [SegBinOp (Aliases GPUMem)]
sbos Names
used_nms = do
([AliasTab]
lutab, [Names]
lu_vars, [Names]
used_nms') <- forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SegBinOp (Aliases GPUMem)
-> LastUseM GPUMem (AliasTab, Names, Names)
helper [SegBinOp (Aliases GPUMem)]
sbos
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Monoid a => [a] -> a
mconcat [AliasTab]
lutab, forall a. Monoid a => [a] -> a
mconcat [Names]
lu_vars, forall a. Monoid a => [a] -> a
mconcat [Names]
used_nms')
where
helper :: SegBinOp (Aliases GPUMem)
-> LastUseM GPUMem (AliasTab, Names, Names)
helper (SegBinOp Commutativity
_ (Lambda [LParam (Aliases GPUMem)]
_ Body (Aliases GPUMem)
body [Type]
_) [SubExp]
neutral Shape
shp) = do
(Names
used_nms', Names
lu_vars) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn [SubExp]
neutral forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn Shape
shp
(AliasTab
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseBody Body (Aliases GPUMem)
body (forall a. Monoid a => a
mempty, Names
used_nms')
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
body_lutab, Names
lu_vars, Names
used_nms'')
lastUseHistOp :: [HistOp (Aliases GPUMem)] -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)
lastUseHistOp :: [HistOp (Aliases GPUMem)]
-> Names -> LastUseM GPUMem (AliasTab, Names, Names)
lastUseHistOp [HistOp (Aliases GPUMem)]
hos Names
used_nms = do
([AliasTab]
lutab, [Names]
lu_vars, [Names]
used_nms') <- forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM HistOp (Aliases GPUMem) -> LastUseM GPUMem (AliasTab, Names, Names)
helper [HistOp (Aliases GPUMem)]
hos
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Monoid a => [a] -> a
mconcat [AliasTab]
lutab, forall a. Monoid a => [a] -> a
mconcat [Names]
lu_vars, forall a. Monoid a => [a] -> a
mconcat [Names]
used_nms')
where
helper :: HistOp (Aliases GPUMem) -> LastUseM GPUMem (AliasTab, Names, Names)
helper (HistOp Shape
shp SubExp
rf [VName]
dest [SubExp]
neutral Shape
shp' (Lambda [LParam (Aliases GPUMem)]
_ Body (Aliases GPUMem)
body [Type]
_)) = do
(Names
used_nms', Names
lu_vars) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn Shape
shp forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn SubExp
rf forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn [VName]
dest forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn [SubExp]
neutral forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn Shape
shp'
(AliasTab
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseBody Body (Aliases GPUMem)
body (forall a. Monoid a => a
mempty, Names
used_nms')
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
body_lutab, Names
lu_vars, Names
used_nms'')
lastUseSeqOp :: Op (Aliases SeqMem) -> Names -> LastUseM SeqMem (LUTabFun, Names, Names)
lastUseSeqOp :: Op (Aliases SeqMem)
-> Names -> LastUseM SeqMem (AliasTab, Names, Names)
lastUseSeqOp (Alloc SubExp
se Space
sp) Names
used_nms = do
let free_in_e :: Names
free_in_e = forall a. FreeIn a => a -> Names
freeIn SubExp
se forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn Space
sp
(Names
used_nms', Names
lu_vars) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms Names
free_in_e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Monoid a => a
mempty, Names
lu_vars, Names
used_nms')
lastUseSeqOp (Inner ()) Names
used_nms = do
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Monoid a => a
mempty, forall a. Monoid a => a
mempty, Names
used_nms)
lastUsedInNames ::
Names ->
Names ->
LastUseM rep (Names, Names)
lastUsedInNames :: forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms Names
new_uses = do
Names
new_uses_with_aliases <- forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure Names
new_uses
[VName]
last_uses <- forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM VName -> StateT AliasTab (Reader (LastUseReader rep)) Bool
isLastUse forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList Names
new_uses
Names
last_uses' <- forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure forall a b. (a -> b) -> a -> b
$ [VName] -> Names
namesFromList [VName]
last_uses
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Names
used_nms forall a. Semigroup a => a -> a -> a
<> Names
new_uses_with_aliases, Names
last_uses')
where
isLastUse :: VName -> StateT AliasTab (Reader (LastUseReader rep)) Bool
isLastUse VName
x = do
Names
with_aliases <- forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure forall a b. (a -> b) -> a -> b
$ VName -> Names
oneName VName
x
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ Names
with_aliases Names -> Names -> Bool
`namesIntersect` Names
used_nms
aliasTransitiveClosure :: Names -> LastUseM rep Names
aliasTransitiveClosure :: forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure Names
args = do
Names
res <- forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall a. Semigroup a => a -> a -> a
(<>) Names
args forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {k} (rep :: k). VName -> LastUseM rep Names
aliasLookup (Names -> [VName]
namesToList Names
args)
if Names
res forall a. Eq a => a -> a -> Bool
== Names
args
then forall (f :: * -> *) a. Applicative f => a -> f a
pure Names
res
else forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure Names
res
updateAliasing ::
AliasesOf dec =>
Names ->
Pat dec ->
LastUseM rep ()
updateAliasing :: forall {k} dec (rep :: k).
AliasesOf dec =>
Names -> Pat dec -> LastUseM rep ()
updateAliasing Names
extra_aliases =
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ forall {k} dec (rep :: k).
AliasesOf dec =>
PatElem dec -> LastUseM rep ()
update forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. Pat dec -> [PatElem dec]
patElems
where
update :: AliasesOf dec => PatElem dec -> LastUseM rep ()
update :: forall {k} dec (rep :: k).
AliasesOf dec =>
PatElem dec -> LastUseM rep ()
update (PatElem VName
name dec
dec) = do
let aliases :: Names
aliases = forall a. AliasesOf a => a -> Names
aliasesOf dec
dec
Names
aliases' <- forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure forall a b. (a -> b) -> a -> b
$ Names
extra_aliases forall a. Semigroup a => a -> a -> a
<> Names
aliases
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
name Names
aliases'