{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeFamilies #-}
module Futhark.Optimise.ArrayShortCircuiting.ArrayCoalescing
( mkCoalsTab,
CoalsTab,
mkCoalsTabGPU,
mkCoalsTabMC,
)
where
import Control.Exception.Base qualified as Exc
import Control.Monad
import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.Function ((&))
import Data.List qualified as L
import Data.List.NonEmpty (NonEmpty (..))
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Sequence (Seq (..))
import Data.Set qualified as S
import Futhark.Analysis.LastUse
import Futhark.Analysis.PrimExp.Convert
import Futhark.IR.Aliases
import Futhark.IR.GPUMem as GPU
import Futhark.IR.MCMem as MC
import Futhark.IR.Mem.IxFun qualified as IxFun
import Futhark.IR.SeqMem
import Futhark.MonadFreshNames
import Futhark.Optimise.ArrayShortCircuiting.DataStructs
import Futhark.Optimise.ArrayShortCircuiting.MemRefAggreg
import Futhark.Optimise.ArrayShortCircuiting.TopdownAnalysis
import Futhark.Util
type Coalesceable rep inner =
( Mem rep inner,
ASTRep rep,
CanBeAliased inner,
AliasableRep rep,
Op rep ~ MemOp inner rep,
HasMemBlock (Aliases rep),
LetDec rep ~ LetDecMem,
TopDownHelper (inner (Aliases rep))
)
type ComputeScalarTable rep op =
ScopeTab rep -> op -> ScalarTableM rep (M.Map VName (PrimExp VName))
newtype ComputeScalarTableOnOp rep = ComputeScalarTableOnOp
{ forall rep.
ComputeScalarTableOnOp rep
-> ComputeScalarTable rep (Op (Aliases rep))
scalarTableOnOp :: ComputeScalarTable rep (Op (Aliases rep))
}
type ScalarTableM rep a = Reader (ComputeScalarTableOnOp rep) a
data ShortCircuitReader rep = ShortCircuitReader
{ forall rep.
ShortCircuitReader rep
-> InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
onOp ::
LUTabFun ->
Pat (VarAliases, LetDecMem) ->
Certs ->
Op (Aliases rep) ->
TopdownEnv rep ->
BotUpEnv ->
ShortCircuitM rep BotUpEnv,
forall rep.
ShortCircuitReader rep
-> InhibitTab
-> TopdownEnv rep
-> ScopeTab rep
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases rep)
-> Maybe [SSPointInfo]
ssPointFromOp ::
LUTabFun ->
TopdownEnv rep ->
ScopeTab rep ->
Pat (VarAliases, LetDecMem) ->
Certs ->
Op (Aliases rep) ->
Maybe [SSPointInfo]
}
newtype ShortCircuitM rep a = ShortCircuitM (ReaderT (ShortCircuitReader rep) (State VNameSource) a)
deriving (forall a b. a -> ShortCircuitM rep b -> ShortCircuitM rep a
forall a b. (a -> b) -> ShortCircuitM rep a -> ShortCircuitM rep b
forall rep a b. a -> ShortCircuitM rep b -> ShortCircuitM rep a
forall rep a b.
(a -> b) -> ShortCircuitM rep a -> ShortCircuitM rep b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> ShortCircuitM rep b -> ShortCircuitM rep a
$c<$ :: forall rep a b. a -> ShortCircuitM rep b -> ShortCircuitM rep a
fmap :: forall a b. (a -> b) -> ShortCircuitM rep a -> ShortCircuitM rep b
$cfmap :: forall rep a b.
(a -> b) -> ShortCircuitM rep a -> ShortCircuitM rep b
Functor, forall rep. Functor (ShortCircuitM rep)
forall a. a -> ShortCircuitM rep a
forall rep a. a -> ShortCircuitM rep a
forall a b.
ShortCircuitM rep a -> ShortCircuitM rep b -> ShortCircuitM rep a
forall a b.
ShortCircuitM rep a -> ShortCircuitM rep b -> ShortCircuitM rep b
forall a b.
ShortCircuitM rep (a -> b)
-> ShortCircuitM rep a -> ShortCircuitM rep b
forall rep a b.
ShortCircuitM rep a -> ShortCircuitM rep b -> ShortCircuitM rep a
forall rep a b.
ShortCircuitM rep a -> ShortCircuitM rep b -> ShortCircuitM rep b
forall rep a b.
ShortCircuitM rep (a -> b)
-> ShortCircuitM rep a -> ShortCircuitM rep b
forall a b c.
(a -> b -> c)
-> ShortCircuitM rep a
-> ShortCircuitM rep b
-> ShortCircuitM rep c
forall rep a b c.
(a -> b -> c)
-> ShortCircuitM rep a
-> ShortCircuitM rep b
-> ShortCircuitM rep 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
<* :: forall a b.
ShortCircuitM rep a -> ShortCircuitM rep b -> ShortCircuitM rep a
$c<* :: forall rep a b.
ShortCircuitM rep a -> ShortCircuitM rep b -> ShortCircuitM rep a
*> :: forall a b.
ShortCircuitM rep a -> ShortCircuitM rep b -> ShortCircuitM rep b
$c*> :: forall rep a b.
ShortCircuitM rep a -> ShortCircuitM rep b -> ShortCircuitM rep b
liftA2 :: forall a b c.
(a -> b -> c)
-> ShortCircuitM rep a
-> ShortCircuitM rep b
-> ShortCircuitM rep c
$cliftA2 :: forall rep a b c.
(a -> b -> c)
-> ShortCircuitM rep a
-> ShortCircuitM rep b
-> ShortCircuitM rep c
<*> :: forall a b.
ShortCircuitM rep (a -> b)
-> ShortCircuitM rep a -> ShortCircuitM rep b
$c<*> :: forall rep a b.
ShortCircuitM rep (a -> b)
-> ShortCircuitM rep a -> ShortCircuitM rep b
pure :: forall a. a -> ShortCircuitM rep a
$cpure :: forall rep a. a -> ShortCircuitM rep a
Applicative, forall rep. Applicative (ShortCircuitM rep)
forall a. a -> ShortCircuitM rep a
forall rep a. a -> ShortCircuitM rep a
forall a b.
ShortCircuitM rep a -> ShortCircuitM rep b -> ShortCircuitM rep b
forall a b.
ShortCircuitM rep a
-> (a -> ShortCircuitM rep b) -> ShortCircuitM rep b
forall rep a b.
ShortCircuitM rep a -> ShortCircuitM rep b -> ShortCircuitM rep b
forall rep a b.
ShortCircuitM rep a
-> (a -> ShortCircuitM rep b) -> ShortCircuitM rep 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
return :: forall a. a -> ShortCircuitM rep a
$creturn :: forall rep a. a -> ShortCircuitM rep a
>> :: forall a b.
ShortCircuitM rep a -> ShortCircuitM rep b -> ShortCircuitM rep b
$c>> :: forall rep a b.
ShortCircuitM rep a -> ShortCircuitM rep b -> ShortCircuitM rep b
>>= :: forall a b.
ShortCircuitM rep a
-> (a -> ShortCircuitM rep b) -> ShortCircuitM rep b
$c>>= :: forall rep a b.
ShortCircuitM rep a
-> (a -> ShortCircuitM rep b) -> ShortCircuitM rep b
Monad, MonadReader (ShortCircuitReader rep), MonadState VNameSource)
instance MonadFreshNames (ShortCircuitM rep) where
putNameSource :: VNameSource -> ShortCircuitM rep ()
putNameSource = forall s (m :: * -> *). MonadState s m => s -> m ()
put
getNameSource :: ShortCircuitM rep VNameSource
getNameSource = forall s (m :: * -> *). MonadState s m => m s
get
emptyTopdownEnv :: TopdownEnv rep
emptyTopdownEnv :: forall rep. TopdownEnv rep
emptyTopdownEnv =
TopdownEnv
{ alloc :: AllocTab
alloc = forall a. Monoid a => a
mempty,
scope :: ScopeTab rep
scope = forall a. Monoid a => a
mempty,
inhibited :: InhibitTab
inhibited = forall a. Monoid a => a
mempty,
v_alias :: VarAliasTab
v_alias = forall a. Monoid a => a
mempty,
m_alias :: InhibitTab
m_alias = forall a. Monoid a => a
mempty,
nonNegatives :: Names
nonNegatives = forall a. Monoid a => a
mempty,
scalarTable :: Map VName (PrimExp VName)
scalarTable = forall a. Monoid a => a
mempty,
knownLessThan :: [(VName, PrimExp VName)]
knownLessThan = forall a. Monoid a => a
mempty,
td_asserts :: [SubExp]
td_asserts = forall a. Monoid a => a
mempty
}
emptyBotUpEnv :: BotUpEnv
emptyBotUpEnv :: BotUpEnv
emptyBotUpEnv =
BotUpEnv
{ scals :: Map VName (PrimExp VName)
scals = forall a. Monoid a => a
mempty,
activeCoals :: CoalsTab
activeCoals = forall a. Monoid a => a
mempty,
successCoals :: CoalsTab
successCoals = forall a. Monoid a => a
mempty,
inhibit :: InhibitTab
inhibit = forall a. Monoid a => a
mempty
}
mkCoalsTab :: (MonadFreshNames m) => Prog (Aliases SeqMem) -> m (M.Map Name CoalsTab)
mkCoalsTab :: forall (m :: * -> *).
MonadFreshNames m =>
Prog (Aliases SeqMem) -> m (Map Name CoalsTab)
mkCoalsTab Prog (Aliases SeqMem)
prog =
forall (m :: * -> *) rep (inner :: * -> *).
(MonadFreshNames m, Coalesceable rep inner) =>
LUTabProg
-> ShortCircuitReader rep
-> ComputeScalarTableOnOp rep
-> Prog (Aliases rep)
-> m (Map Name CoalsTab)
mkCoalsTabProg
(Prog (Aliases SeqMem) -> LUTabProg
lastUseSeqMem Prog (Aliases SeqMem)
prog)
(forall rep.
(InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv)
-> (InhibitTab
-> TopdownEnv rep
-> ScopeTab rep
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases rep)
-> Maybe [SSPointInfo])
-> ShortCircuitReader rep
ShortCircuitReader InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases SeqMem)
-> TopdownEnv SeqMem
-> BotUpEnv
-> ShortCircuitM SeqMem BotUpEnv
shortCircuitSeqMem GenSSPoint SeqMem (Op (Aliases SeqMem))
genSSPointInfoSeqMem)
(forall rep.
ComputeScalarTable rep (Op (Aliases rep))
-> ComputeScalarTableOnOp rep
ComputeScalarTableOnOp forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty)
Prog (Aliases SeqMem)
prog
mkCoalsTabGPU :: (MonadFreshNames m) => Prog (Aliases GPUMem) -> m (M.Map Name CoalsTab)
mkCoalsTabGPU :: forall (m :: * -> *).
MonadFreshNames m =>
Prog (Aliases GPUMem) -> m (Map Name CoalsTab)
mkCoalsTabGPU Prog (Aliases GPUMem)
prog =
forall (m :: * -> *) rep (inner :: * -> *).
(MonadFreshNames m, Coalesceable rep inner) =>
LUTabProg
-> ShortCircuitReader rep
-> ComputeScalarTableOnOp rep
-> Prog (Aliases rep)
-> m (Map Name CoalsTab)
mkCoalsTabProg
(Prog (Aliases GPUMem) -> LUTabProg
lastUseGPUMem Prog (Aliases GPUMem)
prog)
(forall rep.
(InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv)
-> (InhibitTab
-> TopdownEnv rep
-> ScopeTab rep
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases rep)
-> Maybe [SSPointInfo])
-> ShortCircuitReader rep
ShortCircuitReader InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases GPUMem)
-> TopdownEnv GPUMem
-> BotUpEnv
-> ShortCircuitM GPUMem BotUpEnv
shortCircuitGPUMem GenSSPoint GPUMem (Op (Aliases GPUMem))
genSSPointInfoGPUMem)
(forall rep.
ComputeScalarTable rep (Op (Aliases rep))
-> ComputeScalarTableOnOp rep
ComputeScalarTableOnOp (forall rep (inner :: * -> *).
ComputeScalarTable rep (inner (Aliases rep))
-> ComputeScalarTable rep (MemOp inner (Aliases rep))
computeScalarTableMemOp ComputeScalarTable GPUMem (HostOp NoOp (Aliases GPUMem))
computeScalarTableGPUMem))
Prog (Aliases GPUMem)
prog
mkCoalsTabMC :: (MonadFreshNames m) => Prog (Aliases MCMem) -> m (M.Map Name CoalsTab)
mkCoalsTabMC :: forall (m :: * -> *).
MonadFreshNames m =>
Prog (Aliases MCMem) -> m (Map Name CoalsTab)
mkCoalsTabMC Prog (Aliases MCMem)
prog =
forall (m :: * -> *) rep (inner :: * -> *).
(MonadFreshNames m, Coalesceable rep inner) =>
LUTabProg
-> ShortCircuitReader rep
-> ComputeScalarTableOnOp rep
-> Prog (Aliases rep)
-> m (Map Name CoalsTab)
mkCoalsTabProg
(Prog (Aliases MCMem) -> LUTabProg
lastUseMCMem Prog (Aliases MCMem)
prog)
(forall rep.
(InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv)
-> (InhibitTab
-> TopdownEnv rep
-> ScopeTab rep
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases rep)
-> Maybe [SSPointInfo])
-> ShortCircuitReader rep
ShortCircuitReader InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases MCMem)
-> TopdownEnv MCMem
-> BotUpEnv
-> ShortCircuitM MCMem BotUpEnv
shortCircuitMCMem GenSSPoint MCMem (Op (Aliases MCMem))
genSSPointInfoMCMem)
(forall rep.
ComputeScalarTable rep (Op (Aliases rep))
-> ComputeScalarTableOnOp rep
ComputeScalarTableOnOp (forall rep (inner :: * -> *).
ComputeScalarTable rep (inner (Aliases rep))
-> ComputeScalarTable rep (MemOp inner (Aliases rep))
computeScalarTableMemOp ComputeScalarTable MCMem (MCOp NoOp (Aliases MCMem))
computeScalarTableMCMem))
Prog (Aliases MCMem)
prog
mkCoalsTabProg ::
(MonadFreshNames m, Coalesceable rep inner) =>
LUTabProg ->
ShortCircuitReader rep ->
ComputeScalarTableOnOp rep ->
Prog (Aliases rep) ->
m (M.Map Name CoalsTab)
mkCoalsTabProg :: forall (m :: * -> *) rep (inner :: * -> *).
(MonadFreshNames m, Coalesceable rep inner) =>
LUTabProg
-> ShortCircuitReader rep
-> ComputeScalarTableOnOp rep
-> Prog (Aliases rep)
-> m (Map Name CoalsTab)
mkCoalsTabProg (InhibitTab
_, Map Name InhibitTab
lutab_prog) ShortCircuitReader rep
r ComputeScalarTableOnOp rep
computeScalarOnOp =
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM FunDef (Aliases rep) -> m (Name, CoalsTab)
onFun forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Prog rep -> [FunDef rep]
progFuns
where
onFun :: FunDef (Aliases rep) -> m (Name, CoalsTab)
onFun fun :: FunDef (Aliases rep)
fun@(FunDef Maybe EntryPoint
_ Attrs
_ Name
fname [(RetType (Aliases rep), RetAls)]
_ [FParam (Aliases rep)]
fpars Body (Aliases rep)
body) = do
let unique_mems :: AllocTab
unique_mems = [Param FParamMem] -> AllocTab
getUniqueMemFParam [FParam (Aliases rep)]
fpars
lutab :: InhibitTab
lutab = Map Name InhibitTab
lutab_prog forall k a. Ord k => Map k a -> k -> a
M.! Name
fname
scalar_table :: Map VName (PrimExp VName)
scalar_table =
forall r a. Reader r a -> r -> a
runReader
( forall (m :: * -> *) b a.
(Monad m, Monoid b) =>
(a -> m b) -> [a] -> m b
concatMapM
(forall rep (inner :: * -> *).
Coalesceable rep inner =>
ScopeTab rep
-> Stm (Aliases rep)
-> ScalarTableM rep (Map VName (PrimExp VName))
computeScalarTable forall a b. (a -> b) -> a -> b
$ forall rep a. Scoped rep a => a -> Scope rep
scopeOf FunDef (Aliases rep)
fun forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf (forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
body))
(forall rep. Stms rep -> [Stm rep]
stmsToList forall a b. (a -> b) -> a -> b
$ forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
body)
)
ComputeScalarTableOnOp rep
computeScalarOnOp
topenv :: TopdownEnv rep
topenv =
forall rep. TopdownEnv rep
emptyTopdownEnv
{ scope :: ScopeTab rep
scope = forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [FParam (Aliases rep)]
fpars,
alloc :: AllocTab
alloc = AllocTab
unique_mems,
scalarTable :: Map VName (PrimExp VName)
scalarTable = Map VName (PrimExp VName)
scalar_table,
nonNegatives :: Names
nonNegatives = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Param FParamMem -> Names
paramSizes [FParam (Aliases rep)]
fpars
}
ShortCircuitM ReaderT (ShortCircuitReader rep) (State VNameSource) CoalsTab
m = forall rep (inner :: * -> *).
Coalesceable rep inner =>
InhibitTab
-> [Param FParamMem]
-> Body (Aliases rep)
-> TopdownEnv rep
-> ShortCircuitM rep CoalsTab
fixPointCoalesce InhibitTab
lutab [FParam (Aliases rep)]
fpars Body (Aliases rep)
body TopdownEnv rep
topenv
(Name
fname,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource (forall s a. State s a -> s -> (a, s)
runState (forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (ShortCircuitReader rep) (State VNameSource) CoalsTab
m ShortCircuitReader rep
r))
paramSizes :: Param FParamMem -> Names
paramSizes :: Param FParamMem -> Names
paramSizes (Param Attrs
_ VName
_ (MemArray PrimType
_ ShapeBase SubExp
shp Uniqueness
_ MemBind
_)) = forall a. FreeIn a => a -> Names
freeIn ShapeBase SubExp
shp
paramSizes Param FParamMem
_ = forall a. Monoid a => a
mempty
shortCircuitSeqMem :: LUTabFun -> Pat (VarAliases, LetDecMem) -> Certs -> Op (Aliases SeqMem) -> TopdownEnv SeqMem -> BotUpEnv -> ShortCircuitM SeqMem BotUpEnv
shortCircuitSeqMem :: InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases SeqMem)
-> TopdownEnv SeqMem
-> BotUpEnv
-> ShortCircuitM SeqMem BotUpEnv
shortCircuitSeqMem InhibitTab
_ Pat (VarAliases, LetDecMem)
_ Certs
_ Op (Aliases SeqMem)
_ TopdownEnv SeqMem
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure
shortCircuitSegOp ::
(Coalesceable rep inner) =>
(lvl -> Bool) ->
LUTabFun ->
Pat (VarAliases, LetDecMem) ->
Certs ->
SegOp lvl (Aliases rep) ->
TopdownEnv rep ->
BotUpEnv ->
ShortCircuitM rep BotUpEnv
shortCircuitSegOp :: forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
(lvl -> Bool)
-> InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> SegOp lvl (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
shortCircuitSegOp lvl -> Bool
lvlOK InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
pat_certs (SegMap lvl
lvl SegSpace
space [Type]
_ KernelBody (Aliases rep)
kernel_body) TopdownEnv rep
td_env BotUpEnv
bu_env =
forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
Int
-> (lvl -> Bool)
-> lvl
-> InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> SegSpace
-> KernelBody (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
shortCircuitSegOpHelper Int
0 lvl -> Bool
lvlOK lvl
lvl InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
pat_certs SegSpace
space KernelBody (Aliases rep)
kernel_body TopdownEnv rep
td_env BotUpEnv
bu_env
shortCircuitSegOp lvl -> Bool
lvlOK InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
pat_certs (SegRed lvl
lvl SegSpace
space [SegBinOp (Aliases rep)]
binops [Type]
_ KernelBody (Aliases rep)
kernel_body) TopdownEnv rep
td_env BotUpEnv
bu_env =
let to_fail :: CoalsTab
to_fail = forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (\CoalsEntry
entry -> [VName] -> Names
namesFromList (forall k a. Map k a -> [k]
M.keys forall a b. (a -> b) -> a -> b
$ CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
entry) Names -> Names -> Bool
`namesIntersect` forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall a. FreeIn a => a -> Names
freeIn forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda) [SegBinOp (Aliases rep)]
binops) forall a b. (a -> b) -> a -> b
$ BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env
(CoalsTab
active, InhibitTab
inh) =
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env, BotUpEnv -> InhibitTab
inhibit BotUpEnv
bu_env) forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [k]
M.keys CoalsTab
to_fail
bu_env' :: BotUpEnv
bu_env' = BotUpEnv
bu_env {activeCoals :: CoalsTab
activeCoals = CoalsTab
active, inhibit :: InhibitTab
inhibit = InhibitTab
inh}
num_reds :: Int
num_reds = forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
red_ts
in forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
Int
-> (lvl -> Bool)
-> lvl
-> InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> SegSpace
-> KernelBody (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
shortCircuitSegOpHelper Int
num_reds lvl -> Bool
lvlOK lvl
lvl InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
pat_certs SegSpace
space KernelBody (Aliases rep)
kernel_body TopdownEnv rep
td_env BotUpEnv
bu_env'
where
segment_dims :: [SubExp]
segment_dims = forall a. [a] -> [a]
init forall a b. (a -> b) -> a -> b
$ SegSpace -> [SubExp]
segSpaceDims SegSpace
space
red_ts :: [Type]
red_ts = do
SegBinOp (Aliases rep)
op <- [SegBinOp (Aliases rep)]
binops
let shp :: ShapeBase SubExp
shp = forall d. [d] -> ShapeBase d
Shape [SubExp]
segment_dims forall a. Semigroup a => a -> a -> a
<> forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp (Aliases rep)
op
forall a b. (a -> b) -> [a] -> [b]
map (Type -> ShapeBase SubExp -> Type
`arrayOfShape` ShapeBase SubExp
shp) (forall rep. Lambda rep -> [Type]
lambdaReturnType forall a b. (a -> b) -> a -> b
$ forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda SegBinOp (Aliases rep)
op)
shortCircuitSegOp lvl -> Bool
lvlOK InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
pat_certs (SegScan lvl
lvl SegSpace
space [SegBinOp (Aliases rep)]
binops [Type]
_ KernelBody (Aliases rep)
kernel_body) TopdownEnv rep
td_env BotUpEnv
bu_env =
let to_fail :: CoalsTab
to_fail = forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (\CoalsEntry
entry -> [VName] -> Names
namesFromList (forall k a. Map k a -> [k]
M.keys forall a b. (a -> b) -> a -> b
$ CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
entry) Names -> Names -> Bool
`namesIntersect` forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall a. FreeIn a => a -> Names
freeIn forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda) [SegBinOp (Aliases rep)]
binops) forall a b. (a -> b) -> a -> b
$ BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env
(CoalsTab
active, InhibitTab
inh) = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env, BotUpEnv -> InhibitTab
inhibit BotUpEnv
bu_env) forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [k]
M.keys CoalsTab
to_fail
bu_env' :: BotUpEnv
bu_env' = BotUpEnv
bu_env {activeCoals :: CoalsTab
activeCoals = CoalsTab
active, inhibit :: InhibitTab
inhibit = InhibitTab
inh}
in forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
Int
-> (lvl -> Bool)
-> lvl
-> InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> SegSpace
-> KernelBody (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
shortCircuitSegOpHelper Int
0 lvl -> Bool
lvlOK lvl
lvl InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
pat_certs SegSpace
space KernelBody (Aliases rep)
kernel_body TopdownEnv rep
td_env BotUpEnv
bu_env'
shortCircuitSegOp lvl -> Bool
lvlOK InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
pat_certs (SegHist lvl
lvl SegSpace
space [HistOp (Aliases rep)]
histops [Type]
_ KernelBody (Aliases rep)
kernel_body) TopdownEnv rep
td_env BotUpEnv
bu_env = do
let to_fail :: CoalsTab
to_fail = forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (\CoalsEntry
entry -> [VName] -> Names
namesFromList (forall k a. Map k a -> [k]
M.keys forall a b. (a -> b) -> a -> b
$ CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
entry) Names -> Names -> Bool
`namesIntersect` forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall a. FreeIn a => a -> Names
freeIn forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. HistOp rep -> Lambda rep
histOp) [HistOp (Aliases rep)]
histops) forall a b. (a -> b) -> a -> b
$ BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env
(CoalsTab
active, InhibitTab
inh) = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env, BotUpEnv -> InhibitTab
inhibit BotUpEnv
bu_env) forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [k]
M.keys CoalsTab
to_fail
bu_env' :: BotUpEnv
bu_env' = BotUpEnv
bu_env {activeCoals :: CoalsTab
activeCoals = CoalsTab
active, inhibit :: InhibitTab
inhibit = InhibitTab
inh}
BotUpEnv
bu_env'' <- forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
Int
-> (lvl -> Bool)
-> lvl
-> InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> SegSpace
-> KernelBody (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
shortCircuitSegOpHelper Int
0 lvl -> Bool
lvlOK lvl
lvl InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
pat_certs SegSpace
space KernelBody (Aliases rep)
kernel_body TopdownEnv rep
td_env BotUpEnv
bu_env'
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl BotUpEnv -> (PatElem (VarAliases, LetDecMem), VName) -> BotUpEnv
insertHistCoals BotUpEnv
bu_env'' forall a b. (a -> b) -> a -> b
$
forall a b. [a] -> [b] -> [(a, b)]
zip (forall dec. Pat dec -> [PatElem dec]
patElems Pat (VarAliases, LetDecMem)
pat) forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall rep. HistOp rep -> [VName]
histDest [HistOp (Aliases rep)]
histops
where
insertHistCoals :: BotUpEnv -> (PatElem (VarAliases, LetDecMem), VName) -> BotUpEnv
insertHistCoals BotUpEnv
acc (PatElem VName
p (VarAliases, LetDecMem)
_, VName
hist_dest) =
case ( forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
p forall a b. (a -> b) -> a -> b
$ forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env,
forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
hist_dest forall a b. (a -> b) -> a -> b
$ forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env
) of
(Just (MemBlock PrimType
_ ShapeBase SubExp
_ VName
p_mem IxFun
_), Just (MemBlock PrimType
_ ShapeBase SubExp
_ VName
dest_mem IxFun
_)) ->
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
p_mem forall a b. (a -> b) -> a -> b
$ BotUpEnv -> CoalsTab
successCoals BotUpEnv
acc of
Just CoalsEntry
entry ->
let entry' :: CoalsEntry
entry' = CoalsEntry
entry {optdeps :: Map VName VName
optdeps = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
p VName
p_mem forall a b. (a -> b) -> a -> b
$ CoalsEntry -> Map VName VName
optdeps CoalsEntry
entry}
in BotUpEnv
acc
{ successCoals :: CoalsTab
successCoals = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
p_mem CoalsEntry
entry' forall a b. (a -> b) -> a -> b
$ BotUpEnv -> CoalsTab
successCoals BotUpEnv
acc,
activeCoals :: CoalsTab
activeCoals = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
dest_mem CoalsEntry
entry forall a b. (a -> b) -> a -> b
$ BotUpEnv -> CoalsTab
activeCoals BotUpEnv
acc
}
Maybe CoalsEntry
Nothing -> BotUpEnv
acc
(Maybe ArrayMemBound, Maybe ArrayMemBound)
_ -> BotUpEnv
acc
shortCircuitGPUMem ::
LUTabFun ->
Pat (VarAliases, LetDecMem) ->
Certs ->
Op (Aliases GPUMem) ->
TopdownEnv GPUMem ->
BotUpEnv ->
ShortCircuitM GPUMem BotUpEnv
shortCircuitGPUMem :: InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases GPUMem)
-> TopdownEnv GPUMem
-> BotUpEnv
-> ShortCircuitM GPUMem BotUpEnv
shortCircuitGPUMem InhibitTab
_ Pat (VarAliases, LetDecMem)
_ Certs
_ (Alloc SubExp
_ Space
_) TopdownEnv GPUMem
_ BotUpEnv
bu_env = forall (f :: * -> *) a. Applicative f => a -> f a
pure BotUpEnv
bu_env
shortCircuitGPUMem InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
certs (Inner (GPU.SegOp SegOp SegLevel (Aliases GPUMem)
op)) TopdownEnv GPUMem
td_env BotUpEnv
bu_env =
forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
(lvl -> Bool)
-> InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> SegOp lvl (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
shortCircuitSegOp SegLevel -> Bool
isSegThread InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
certs SegOp SegLevel (Aliases GPUMem)
op TopdownEnv GPUMem
td_env BotUpEnv
bu_env
shortCircuitGPUMem InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
certs (Inner (GPU.GPUBody [Type]
_ Body (Aliases GPUMem)
body)) TopdownEnv GPUMem
td_env BotUpEnv
bu_env = do
VName
fresh1 <- forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString String
"gpubody"
VName
fresh2 <- forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString String
"gpubody"
forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
Int
-> (lvl -> Bool)
-> lvl
-> InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> SegSpace
-> KernelBody (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
shortCircuitSegOpHelper
Int
0
SegLevel -> Bool
isSegThread
( SegVirt -> Maybe KernelGrid -> SegLevel
GPU.SegThread SegVirt
GPU.SegNoVirt forall a b. (a -> b) -> a -> b
$
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
Count NumGroups SubExp -> Count GroupSize SubExp -> KernelGrid
GPU.KernelGrid
(forall {k} (u :: k) e. e -> Count u e
GPU.Count forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
IntValue forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value Int64
1)
(forall {k} (u :: k) e. e -> Count u e
GPU.Count forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
IntValue forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value Int64
1)
)
InhibitTab
lutab
Pat (VarAliases, LetDecMem)
pat
Certs
certs
(VName -> [(VName, SubExp)] -> SegSpace
SegSpace VName
fresh1 [(VName
fresh2, PrimValue -> SubExp
Constant forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
IntValue forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value Int64
1)])
(Body (Aliases GPUMem) -> KernelBody (Aliases GPUMem)
bodyToKernelBody Body (Aliases GPUMem)
body)
TopdownEnv GPUMem
td_env
BotUpEnv
bu_env
shortCircuitGPUMem InhibitTab
_ Pat (VarAliases, LetDecMem)
_ Certs
_ (Inner (GPU.SizeOp SizeOp
_)) TopdownEnv GPUMem
_ BotUpEnv
bu_env = forall (f :: * -> *) a. Applicative f => a -> f a
pure BotUpEnv
bu_env
shortCircuitGPUMem InhibitTab
_ Pat (VarAliases, LetDecMem)
_ Certs
_ (Inner (GPU.OtherOp NoOp (Aliases GPUMem)
NoOp)) TopdownEnv GPUMem
_ BotUpEnv
bu_env = forall (f :: * -> *) a. Applicative f => a -> f a
pure BotUpEnv
bu_env
shortCircuitMCMem ::
LUTabFun ->
Pat (VarAliases, LetDecMem) ->
Certs ->
Op (Aliases MCMem) ->
TopdownEnv MCMem ->
BotUpEnv ->
ShortCircuitM MCMem BotUpEnv
shortCircuitMCMem :: InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases MCMem)
-> TopdownEnv MCMem
-> BotUpEnv
-> ShortCircuitM MCMem BotUpEnv
shortCircuitMCMem InhibitTab
_ Pat (VarAliases, LetDecMem)
_ Certs
_ (Alloc SubExp
_ Space
_) TopdownEnv MCMem
_ BotUpEnv
bu_env = forall (f :: * -> *) a. Applicative f => a -> f a
pure BotUpEnv
bu_env
shortCircuitMCMem InhibitTab
_ Pat (VarAliases, LetDecMem)
_ Certs
_ (Inner (MC.OtherOp NoOp (Aliases MCMem)
NoOp)) TopdownEnv MCMem
_ BotUpEnv
bu_env = forall (f :: * -> *) a. Applicative f => a -> f a
pure BotUpEnv
bu_env
shortCircuitMCMem InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
certs (Inner (MC.ParOp (Just SegOp () (Aliases MCMem)
par_op) SegOp () (Aliases MCMem)
op)) TopdownEnv MCMem
td_env BotUpEnv
bu_env =
forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
(lvl -> Bool)
-> InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> SegOp lvl (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
shortCircuitSegOp (forall a b. a -> b -> a
const Bool
True) InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
certs SegOp () (Aliases MCMem)
par_op TopdownEnv MCMem
td_env BotUpEnv
bu_env
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
(lvl -> Bool)
-> InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> SegOp lvl (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
shortCircuitSegOp (forall a b. a -> b -> a
const Bool
True) InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
certs SegOp () (Aliases MCMem)
op TopdownEnv MCMem
td_env
shortCircuitMCMem InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
certs (Inner (MC.ParOp Maybe (SegOp () (Aliases MCMem))
Nothing SegOp () (Aliases MCMem)
op)) TopdownEnv MCMem
td_env BotUpEnv
bu_env =
forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
(lvl -> Bool)
-> InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> SegOp lvl (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
shortCircuitSegOp (forall a b. a -> b -> a
const Bool
True) InhibitTab
lutab Pat (VarAliases, LetDecMem)
pat Certs
certs SegOp () (Aliases MCMem)
op TopdownEnv MCMem
td_env BotUpEnv
bu_env
dropLastSegSpace :: SegSpace -> SegSpace
dropLastSegSpace :: SegSpace -> SegSpace
dropLastSegSpace SegSpace
space = SegSpace
space {unSegSpace :: [(VName, SubExp)]
unSegSpace = forall a. [a] -> [a]
init forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space}
isSegThread :: GPU.SegLevel -> Bool
isSegThread :: SegLevel -> Bool
isSegThread GPU.SegThread {} = Bool
True
isSegThread SegLevel
_ = Bool
False
threadSlice :: SegSpace -> KernelResult -> Maybe (Slice (TPrimExp Int64 VName))
threadSlice :: SegSpace -> KernelResult -> Maybe (Slice (TPrimExp Int64 VName))
threadSlice SegSpace
space Returns {} =
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
forall d. [DimIndex d] -> Slice d
Slice forall a b. (a -> b) -> a -> b
$
forall a b. (a -> b) -> [a] -> [b]
map (forall d. d -> DimIndex d
DimFix forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (t :: k) v. PrimExp v -> TPrimExp t v
TPrimExp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip forall v. v -> PrimType -> PrimExp v
LeafExp (IntType -> PrimType
IntType IntType
Int64) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall a b. (a -> b) -> a -> b
$
SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space
threadSlice SegSpace
space (RegTileReturns Certs
_ [(SubExp, SubExp, SubExp)]
dims VName
_) =
forall a. a -> Maybe a
Just
forall a b. (a -> b) -> a -> b
$ forall d. [DimIndex d] -> Slice d
Slice
forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
( \(SubExp
_, SubExp
block_tile_size0, SubExp
reg_tile_size0) (VName
x0, SubExp
_) ->
let x :: TPrimExp Int64 VName
x = SubExp -> TPrimExp Int64 VName
pe64 forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
x0
block_tile_size :: TPrimExp Int64 VName
block_tile_size = SubExp -> TPrimExp Int64 VName
pe64 SubExp
block_tile_size0
reg_tile_size :: TPrimExp Int64 VName
reg_tile_size = SubExp -> TPrimExp Int64 VName
pe64 SubExp
reg_tile_size0
in forall d. d -> d -> d -> DimIndex d
DimSlice (TPrimExp Int64 VName
x forall a. Num a => a -> a -> a
* TPrimExp Int64 VName
block_tile_size forall a. Num a => a -> a -> a
* TPrimExp Int64 VName
reg_tile_size) (TPrimExp Int64 VName
block_tile_size forall a. Num a => a -> a -> a
* TPrimExp Int64 VName
reg_tile_size) TPrimExp Int64 VName
1
)
[(SubExp, SubExp, SubExp)]
dims
forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space
threadSlice SegSpace
_ KernelResult
_ = forall a. Maybe a
Nothing
bodyToKernelBody :: Body (Aliases GPUMem) -> KernelBody (Aliases GPUMem)
bodyToKernelBody :: Body (Aliases GPUMem) -> KernelBody (Aliases GPUMem)
bodyToKernelBody (Body BodyDec (Aliases GPUMem)
dec Stms (Aliases GPUMem)
stms Result
res) =
forall rep.
BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
KernelBody BodyDec (Aliases GPUMem)
dec Stms (Aliases GPUMem)
stms forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (\(SubExpRes Certs
cert SubExp
subexps) -> ResultManifest -> Certs -> SubExp -> KernelResult
Returns ResultManifest
ResultNoSimplify Certs
cert SubExp
subexps) Result
res
shortCircuitSegOpHelper ::
(Coalesceable rep inner) =>
Int ->
(lvl -> Bool) ->
lvl ->
LUTabFun ->
Pat (VarAliases, LetDecMem) ->
Certs ->
SegSpace ->
KernelBody (Aliases rep) ->
TopdownEnv rep ->
BotUpEnv ->
ShortCircuitM rep BotUpEnv
shortCircuitSegOpHelper :: forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
Int
-> (lvl -> Bool)
-> lvl
-> InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> SegSpace
-> KernelBody (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
shortCircuitSegOpHelper Int
num_reds lvl -> Bool
lvlOK lvl
lvl InhibitTab
lutab pat :: Pat (VarAliases, LetDecMem)
pat@(Pat [PatElem (VarAliases, LetDecMem)]
ps0) Certs
pat_certs SegSpace
space0 KernelBody (Aliases rep)
kernel_body TopdownEnv rep
td_env BotUpEnv
bu_env = do
let ps_space_and_res :: [(PatElem (VarAliases, LetDecMem), SegSpace, KernelResult)]
ps_space_and_res =
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [PatElem (VarAliases, LetDecMem)]
ps0 (forall a. Int -> a -> [a]
replicate Int
num_reds (SegSpace -> SegSpace
dropLastSegSpace SegSpace
space0) forall a. Semigroup a => a -> a -> a
<> forall a. a -> [a]
repeat SegSpace
space0) forall a b. (a -> b) -> a -> b
$
forall rep. KernelBody rep -> [KernelResult]
kernelBodyResult KernelBody (Aliases rep)
kernel_body
let (CoalsTab
actv0, InhibitTab
inhibit0) =
forall rep.
HasMemBlock (Aliases rep) =>
CoalsTab
-> InhibitTab
-> Map VName (PrimExp VName)
-> TopdownEnv rep
-> [PatElem (VarAliases, LetDecMem)]
-> (CoalsTab, InhibitTab)
filterSafetyCond2and5
(BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env)
(BotUpEnv -> InhibitTab
inhibit BotUpEnv
bu_env)
(BotUpEnv -> Map VName (PrimExp VName)
scals BotUpEnv
bu_env)
TopdownEnv rep
td_env
(forall dec. Pat dec -> [PatElem dec]
patElems Pat (VarAliases, LetDecMem)
pat)
(CoalsTab
actv_return, InhibitTab
inhibit_return) =
if Int
num_reds forall a. Ord a => a -> a -> Bool
> Int
0
then (CoalsTab
actv0, InhibitTab
inhibit0)
else forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
(lvl -> Bool)
-> lvl
-> TopdownEnv rep
-> KernelBody (Aliases rep)
-> Certs
-> (CoalsTab, InhibitTab)
-> (PatElem (VarAliases, LetDecMem), SegSpace, KernelResult)
-> (CoalsTab, InhibitTab)
makeSegMapCoals lvl -> Bool
lvlOK lvl
lvl TopdownEnv rep
td_env KernelBody (Aliases rep)
kernel_body Certs
pat_certs) (CoalsTab
actv0, InhibitTab
inhibit0) [(PatElem (VarAliases, LetDecMem), SegSpace, KernelResult)]
ps_space_and_res
let actv0' :: CoalsTab
actv0' = forall a b k. (a -> b) -> Map k a -> Map k b
M.map (\CoalsEntry
etry -> CoalsEntry
etry {memrefs :: MemRefs
memrefs = forall a. Monoid a => a
mempty}) forall a b. (a -> b) -> a -> b
$ CoalsTab
actv0 forall a. Semigroup a => a -> a -> a
<> CoalsTab
actv_return
BotUpEnv
bu_env' <-
forall rep (inner :: * -> *).
Coalesceable rep inner =>
InhibitTab
-> Stms (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
mkCoalsTabStms InhibitTab
lutab (forall rep. KernelBody rep -> Stms rep
kernelBodyStms KernelBody (Aliases rep)
kernel_body) TopdownEnv rep
td_env forall a b. (a -> b) -> a -> b
$
BotUpEnv
bu_env {activeCoals :: CoalsTab
activeCoals = CoalsTab
actv0', inhibit :: InhibitTab
inhibit = InhibitTab
inhibit_return}
let actv_coals_after :: CoalsTab
actv_coals_after =
forall k a b. (k -> a -> b) -> Map k a -> Map k b
M.mapWithKey
( \VName
k CoalsEntry
etry ->
CoalsEntry
etry
{ memrefs :: MemRefs
memrefs = CoalsEntry -> MemRefs
memrefs CoalsEntry
etry forall a. Semigroup a => a -> a -> a
<> forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty CoalsEntry -> MemRefs
memrefs (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
k forall a b. (a -> b) -> a -> b
$ CoalsTab
actv0 forall a. Semigroup a => a -> a -> a
<> CoalsTab
actv_return)
}
)
forall a b. (a -> b) -> a -> b
$ BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env'
let checkPartialOverlap :: BotUpEnv -> (VName, CoalsEntry) -> ShortCircuitM rep BotUpEnv
checkPartialOverlap BotUpEnv
bu_env_f (VName
k, CoalsEntry
entry) = do
let sliceThreadAccess :: (PatElem (VarAliases, LetDecMem), SegSpace, KernelResult)
-> AccessSummary
sliceThreadAccess (PatElem (VarAliases, LetDecMem)
p, SegSpace
space, KernelResult
res) =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall dec. PatElem dec -> VName
patElemName PatElem (VarAliases, LetDecMem)
p) forall a b. (a -> b) -> a -> b
$ CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
entry of
Just (Coalesced CoalescedKind
_ (MemBlock PrimType
_ ShapeBase SubExp
_ VName
_ IxFun
ixf) FreeVarSubsts
_) ->
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
AccessSummary
Undeterminable
( IxFun -> AccessSummary
ixfunToAccessSummary
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall num.
(Eq num, IntegralExp num) =>
IxFun num -> Slice num -> IxFun num
IxFun.slice IxFun
ixf
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TPrimExp Int64 VName]
-> Slice (TPrimExp Int64 VName) -> Slice (TPrimExp Int64 VName)
fullSlice (forall num. (Eq num, IntegralExp num) => IxFun num -> Shape num
IxFun.shape IxFun
ixf)
)
forall a b. (a -> b) -> a -> b
$ SegSpace -> KernelResult -> Maybe (Slice (TPrimExp Int64 VName))
threadSlice SegSpace
space KernelResult
res
Maybe Coalesced
Nothing -> forall a. Monoid a => a
mempty
thread_writes :: AccessSummary
thread_writes = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (PatElem (VarAliases, LetDecMem), SegSpace, KernelResult)
-> AccessSummary
sliceThreadAccess [(PatElem (VarAliases, LetDecMem), SegSpace, KernelResult)]
ps_space_and_res
source_writes :: AccessSummary
source_writes = MemRefs -> AccessSummary
srcwrts (CoalsEntry -> MemRefs
memrefs CoalsEntry
entry) forall a. Semigroup a => a -> a -> a
<> AccessSummary
thread_writes
AccessSummary
destination_uses <-
case MemRefs -> AccessSummary
dstrefs (CoalsEntry -> MemRefs
memrefs CoalsEntry
entry)
AccessSummary -> AccessSummary -> AccessSummary
`accessSubtract` MemRefs -> AccessSummary
dstrefs (forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty CoalsEntry -> MemRefs
memrefs forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
k forall a b. (a -> b) -> a -> b
$ BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env) of
Set Set LmadRef
s ->
forall (m :: * -> *) b a.
(Monad m, Monoid b) =>
(a -> m b) -> [a] -> m b
concatMapM
(forall (m :: * -> *).
MonadFreshNames m =>
Map VName (PrimExp VName)
-> [(VName, SubExp)] -> LmadRef -> m AccessSummary
aggSummaryMapPartial (forall rep. TopdownEnv rep -> Map VName (PrimExp VName)
scalarTable TopdownEnv rep
td_env) forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space0)
(forall a. Set a -> [a]
S.toList Set LmadRef
s)
AccessSummary
Undeterminable -> forall (f :: * -> *) a. Applicative f => a -> f a
pure AccessSummary
Undeterminable
if CoalsEntry -> VName
dstmem CoalsEntry
entry forall k a. Ord k => k -> Map k a -> Bool
`M.member` forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env
Bool -> Bool -> Bool
&& forall rep.
AliasableRep rep =>
TopdownEnv rep -> AccessSummary -> AccessSummary -> Bool
noMemOverlap TopdownEnv rep
td_env AccessSummary
destination_uses AccessSummary
source_writes
then forall (f :: * -> *) a. Applicative f => a -> f a
pure BotUpEnv
bu_env_f
else do
let (CoalsTab
ac, InhibitTab
inh) = (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env_f, BotUpEnv -> InhibitTab
inhibit BotUpEnv
bu_env_f) VName
k
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ BotUpEnv
bu_env_f {activeCoals :: CoalsTab
activeCoals = CoalsTab
ac, inhibit :: InhibitTab
inhibit = InhibitTab
inh}
BotUpEnv
bu_env'' <-
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
BotUpEnv -> (VName, CoalsEntry) -> ShortCircuitM rep BotUpEnv
checkPartialOverlap
(BotUpEnv
bu_env' {activeCoals :: CoalsTab
activeCoals = CoalsTab
actv_coals_after})
forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList CoalsTab
actv_coals_after
let updateMemRefs :: CoalsEntry -> ShortCircuitM rep CoalsEntry
updateMemRefs CoalsEntry
entry = do
AccessSummary
wrts <- forall (m :: * -> *).
MonadFreshNames m =>
Map VName (PrimExp VName)
-> [(VName, SubExp)] -> AccessSummary -> m AccessSummary
aggSummaryMapTotal (forall rep. TopdownEnv rep -> Map VName (PrimExp VName)
scalarTable TopdownEnv rep
td_env) (SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space0) forall a b. (a -> b) -> a -> b
$ MemRefs -> AccessSummary
srcwrts forall a b. (a -> b) -> a -> b
$ CoalsEntry -> MemRefs
memrefs CoalsEntry
entry
AccessSummary
uses <- forall (m :: * -> *).
MonadFreshNames m =>
Map VName (PrimExp VName)
-> [(VName, SubExp)] -> AccessSummary -> m AccessSummary
aggSummaryMapTotal (forall rep. TopdownEnv rep -> Map VName (PrimExp VName)
scalarTable TopdownEnv rep
td_env) (SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space0) forall a b. (a -> b) -> a -> b
$ MemRefs -> AccessSummary
dstrefs forall a b. (a -> b) -> a -> b
$ CoalsEntry -> MemRefs
memrefs CoalsEntry
entry
let uses' :: AccessSummary
uses' =
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap
( \case
PatElem VName
_ (VarAliases
_, MemArray PrimType
_ ShapeBase SubExp
_ NoUniqueness
_ (ArrayIn VName
p_mem IxFun
p_ixf))
| VName
p_mem VName -> Names -> Bool
`nameIn` CoalsEntry -> Names
alsmem CoalsEntry
entry ->
IxFun -> AccessSummary
ixfunToAccessSummary IxFun
p_ixf
PatElem (VarAliases, LetDecMem)
_ -> forall a. Monoid a => a
mempty
)
[PatElem (VarAliases, LetDecMem)]
ps0
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ CoalsEntry
entry {memrefs :: MemRefs
memrefs = AccessSummary -> AccessSummary -> MemRefs
MemRefs (AccessSummary
uses forall a. Semigroup a => a -> a -> a
<> AccessSummary
uses') AccessSummary
wrts}
CoalsTab
actv <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM CoalsEntry -> ShortCircuitM rep CoalsEntry
updateMemRefs forall a b. (a -> b) -> a -> b
$ BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env''
let bu_env''' :: BotUpEnv
bu_env''' = BotUpEnv
bu_env'' {activeCoals :: CoalsTab
activeCoals = CoalsTab
actv}
let mergee_writes :: [(PatElem (VarAliases, LetDecMem), (VName, VName, IxFun))]
mergee_writes =
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
( \(PatElem (VarAliases, LetDecMem)
p, SegSpace
_, KernelResult
_) ->
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (PatElem (VarAliases, LetDecMem)
p,) forall a b. (a -> b) -> a -> b
$
forall rep.
HasMemBlock (Aliases rep) =>
TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, IxFun)
getDirAliasedIxfn' TopdownEnv rep
td_env (BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env''') forall a b. (a -> b) -> a -> b
$
forall dec. PatElem dec -> VName
patElemName PatElem (VarAliases, LetDecMem)
p
)
[(PatElem (VarAliases, LetDecMem), SegSpace, KernelResult)]
ps_space_and_res
let checkMergeeOverlap :: BotUpEnv
-> (PatElem (VarAliases, LetDecMem), (VName, VName, IxFun))
-> ShortCircuitM rep BotUpEnv
checkMergeeOverlap BotUpEnv
bu_env_f (PatElem (VarAliases, LetDecMem)
p, (VName
m_b, VName
_, IxFun
ixf)) =
let as :: AccessSummary
as = IxFun -> AccessSummary
ixfunToAccessSummary IxFun
ixf
in
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_b forall a b. (a -> b) -> a -> b
$ BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env of
Just CoalsEntry
coal_entry -> do
let mrefs :: MemRefs
mrefs =
CoalsEntry -> MemRefs
memrefs CoalsEntry
coal_entry
res :: Bool
res = forall rep.
AliasableRep rep =>
TopdownEnv rep -> AccessSummary -> AccessSummary -> Bool
noMemOverlap TopdownEnv rep
td_env AccessSummary
as forall a b. (a -> b) -> a -> b
$ MemRefs -> AccessSummary
dstrefs MemRefs
mrefs
fail_res :: BotUpEnv
fail_res =
let (CoalsTab
ac, InhibitTab
inh) = (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env_f, BotUpEnv -> InhibitTab
inhibit BotUpEnv
bu_env_f) VName
m_b
in BotUpEnv
bu_env_f {activeCoals :: CoalsTab
activeCoals = CoalsTab
ac, inhibit :: InhibitTab
inhibit = InhibitTab
inh}
if Bool
res
then case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall dec. PatElem dec -> VName
patElemName PatElem (VarAliases, LetDecMem)
p) forall a b. (a -> b) -> a -> b
$ CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
coal_entry of
Maybe Coalesced
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure BotUpEnv
bu_env_f
Just (Coalesced CoalescedKind
knd mbd :: ArrayMemBound
mbd@(MemBlock PrimType
_ ShapeBase SubExp
_ VName
_ IxFun
ixfn) FreeVarSubsts
_) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
case forall a rep.
FreeIn a =>
ScopeTab rep
-> Map VName (PrimExp VName) -> a -> Maybe FreeVarSubsts
freeVarSubstitutions (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) (forall rep. TopdownEnv rep -> Map VName (PrimExp VName)
scalarTable TopdownEnv rep
td_env) IxFun
ixfn of
Just FreeVarSubsts
fv_subst ->
let entry :: CoalsEntry
entry =
CoalsEntry
coal_entry
{ vartab :: Map VName Coalesced
vartab =
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert
(forall dec. PatElem dec -> VName
patElemName PatElem (VarAliases, LetDecMem)
p)
(CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced CoalescedKind
knd ArrayMemBound
mbd FreeVarSubsts
fv_subst)
(CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
coal_entry)
}
(CoalsTab
ac, CoalsTab
suc) =
(CoalsTab, CoalsTab) -> VName -> CoalsEntry -> (CoalsTab, CoalsTab)
markSuccessCoal (BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env_f, BotUpEnv -> CoalsTab
successCoals BotUpEnv
bu_env_f) VName
m_b CoalsEntry
entry
in BotUpEnv
bu_env_f {activeCoals :: CoalsTab
activeCoals = CoalsTab
ac, successCoals :: CoalsTab
successCoals = CoalsTab
suc}
Maybe FreeVarSubsts
Nothing ->
BotUpEnv
fail_res
else forall (f :: * -> *) a. Applicative f => a -> f a
pure BotUpEnv
fail_res
Maybe CoalsEntry
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure BotUpEnv
bu_env_f
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM BotUpEnv
-> (PatElem (VarAliases, LetDecMem), (VName, VName, IxFun))
-> ShortCircuitM rep BotUpEnv
checkMergeeOverlap BotUpEnv
bu_env''' [(PatElem (VarAliases, LetDecMem), (VName, VName, IxFun))]
mergee_writes
makeSegMapCoals ::
(Coalesceable rep inner) =>
(lvl -> Bool) ->
lvl ->
TopdownEnv rep ->
KernelBody (Aliases rep) ->
Certs ->
(CoalsTab, InhibitTab) ->
(PatElem (VarAliases, LetDecMem), SegSpace, KernelResult) ->
(CoalsTab, InhibitTab)
makeSegMapCoals :: forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
(lvl -> Bool)
-> lvl
-> TopdownEnv rep
-> KernelBody (Aliases rep)
-> Certs
-> (CoalsTab, InhibitTab)
-> (PatElem (VarAliases, LetDecMem), SegSpace, KernelResult)
-> (CoalsTab, InhibitTab)
makeSegMapCoals lvl -> Bool
lvlOK lvl
lvl TopdownEnv rep
td_env KernelBody (Aliases rep)
kernel_body Certs
pat_certs (CoalsTab
active, InhibitTab
inhb) (PatElem VName
pat_name (VarAliases
_, MemArray PrimType
_ ShapeBase SubExp
_ NoUniqueness
_ (ArrayIn VName
pat_mem IxFun
pat_ixf)), SegSpace
space, Returns ResultManifest
_ Certs
_ (Var VName
return_name))
| Just (MemBlock PrimType
tp ShapeBase SubExp
return_shp VName
return_mem IxFun
_) <-
forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
return_name forall a b. (a -> b) -> a -> b
$ forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf (forall rep. KernelBody rep -> Stms rep
kernelBodyStms KernelBody (Aliases rep)
kernel_body),
lvl -> Bool
lvlOK lvl
lvl,
MemMem Space
pat_space <- forall r a. Reader r a -> r -> a
runReader (forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Mem rep inner) =>
VName -> m LetDecMem
lookupMemInfo VName
pat_mem) forall a b. (a -> b) -> a -> b
$ forall rep. Scope (Aliases rep) -> Scope rep
removeScopeAliases forall a b. (a -> b) -> a -> b
$ forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env,
MemMem Space
return_space <-
forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf (forall rep. KernelBody rep -> Stms rep
kernelBodyStms KernelBody (Aliases rep)
kernel_body) forall a. Semigroup a => a -> a -> a
<> forall rep. SegSpace -> Scope rep
scopeOfSegSpace SegSpace
space
forall a b. a -> (a -> b) -> b
& forall rep. Scope (Aliases rep) -> Scope rep
removeScopeAliases
forall a b. a -> (a -> b) -> b
& forall r a. Reader r a -> r -> a
runReader (forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Mem rep inner) =>
VName -> m LetDecMem
lookupMemInfo VName
return_mem),
Space
pat_space forall a. Eq a => a -> a -> Bool
== Space
return_space =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
pat_mem CoalsTab
active of
Maybe CoalsEntry
Nothing ->
case ( forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (VName
pat_mem `nameIn`) (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
return_mem InhibitTab
inhb),
CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced
CoalescedKind
InPlaceCoal
(PrimType -> ShapeBase SubExp -> VName -> IxFun -> ArrayMemBound
MemBlock PrimType
tp ShapeBase SubExp
return_shp VName
pat_mem forall a b. (a -> b) -> a -> b
$ IxFun -> IxFun
resultSlice IxFun
pat_ixf)
forall a. Monoid a => a
mempty
forall a b. a -> (a -> b) -> b
& forall k a. k -> a -> Map k a
M.singleton VName
return_name
forall a b. a -> (a -> b) -> b
& forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall rep.
HasMemBlock (Aliases rep) =>
TopdownEnv rep
-> Map VName Coalesced -> VName -> Maybe (Map VName Coalesced)
addInvAliasesVarTab TopdownEnv rep
td_env) VName
return_name
) of
(Bool
False, Just Map VName Coalesced
vtab) ->
( CoalsTab
active
forall a. Semigroup a => a -> a -> a
<> forall k a. k -> a -> Map k a
M.singleton
VName
return_mem
(VName
-> IxFun
-> Names
-> Map VName Coalesced
-> Map VName VName
-> MemRefs
-> Certs
-> CoalsEntry
CoalsEntry VName
pat_mem IxFun
pat_ixf (VName -> Names
oneName VName
pat_mem) Map VName Coalesced
vtab forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty Certs
pat_certs),
InhibitTab
inhb
)
(Bool, Maybe (Map VName Coalesced))
_ -> (CoalsTab
active, InhibitTab
inhb)
Just CoalsEntry
trans ->
case ( forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (CoalsEntry -> VName
dstmem CoalsEntry
trans `nameIn`) forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
return_mem InhibitTab
inhb,
let Coalesced CoalescedKind
_ (MemBlock PrimType
_ ShapeBase SubExp
_ VName
trans_mem IxFun
trans_ixf) FreeVarSubsts
_ =
forall a. a -> Maybe a -> a
fromMaybe (forall a. HasCallStack => String -> a
error String
"Impossible") forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
pat_name forall a b. (a -> b) -> a -> b
$ CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
trans
in CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced
CoalescedKind
TransitiveCoal
(PrimType -> ShapeBase SubExp -> VName -> IxFun -> ArrayMemBound
MemBlock PrimType
tp ShapeBase SubExp
return_shp VName
trans_mem forall a b. (a -> b) -> a -> b
$ IxFun -> IxFun
resultSlice IxFun
trans_ixf)
forall a. Monoid a => a
mempty
forall a b. a -> (a -> b) -> b
& forall k a. k -> a -> Map k a
M.singleton VName
return_name
forall a b. a -> (a -> b) -> b
& forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall rep.
HasMemBlock (Aliases rep) =>
TopdownEnv rep
-> Map VName Coalesced -> VName -> Maybe (Map VName Coalesced)
addInvAliasesVarTab TopdownEnv rep
td_env) VName
return_name
) of
(Bool
False, Just Map VName Coalesced
vtab) ->
let opts :: Map VName VName
opts =
if CoalsEntry -> VName
dstmem CoalsEntry
trans forall a. Eq a => a -> a -> Bool
== VName
pat_mem
then forall a. Monoid a => a
mempty
else forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
pat_name VName
pat_mem forall a b. (a -> b) -> a -> b
$ CoalsEntry -> Map VName VName
optdeps CoalsEntry
trans
in ( forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert
VName
return_mem
( VName
-> IxFun
-> Names
-> Map VName Coalesced
-> Map VName VName
-> MemRefs
-> Certs
-> CoalsEntry
CoalsEntry
(CoalsEntry -> VName
dstmem CoalsEntry
trans)
(CoalsEntry -> IxFun
dstind CoalsEntry
trans)
(VName -> Names
oneName VName
pat_mem forall a. Semigroup a => a -> a -> a
<> CoalsEntry -> Names
alsmem CoalsEntry
trans)
Map VName Coalesced
vtab
Map VName VName
opts
forall a. Monoid a => a
mempty
(CoalsEntry -> Certs
certs CoalsEntry
trans forall a. Semigroup a => a -> a -> a
<> Certs
pat_certs)
)
CoalsTab
active,
InhibitTab
inhb
)
(Bool, Maybe (Map VName Coalesced))
_ -> (CoalsTab
active, InhibitTab
inhb)
where
thread_slice :: Slice (TPrimExp Int64 VName)
thread_slice =
SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space
forall a b. a -> (a -> b) -> b
& forall a b. (a -> b) -> [a] -> [b]
map (forall d. d -> DimIndex d
DimFix forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (t :: k) v. PrimExp v -> TPrimExp t v
TPrimExp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip forall v. v -> PrimType -> PrimExp v
LeafExp (IntType -> PrimType
IntType IntType
Int64) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst)
forall a b. a -> (a -> b) -> b
& forall d. [DimIndex d] -> Slice d
Slice
resultSlice :: IxFun -> IxFun
resultSlice IxFun
ixf = forall num.
(Eq num, IntegralExp num) =>
IxFun num -> Slice num -> IxFun num
IxFun.slice IxFun
ixf forall a b. (a -> b) -> a -> b
$ [TPrimExp Int64 VName]
-> Slice (TPrimExp Int64 VName) -> Slice (TPrimExp Int64 VName)
fullSlice (forall num. (Eq num, IntegralExp num) => IxFun num -> Shape num
IxFun.shape IxFun
ixf) Slice (TPrimExp Int64 VName)
thread_slice
makeSegMapCoals lvl -> Bool
_ lvl
_ TopdownEnv rep
td_env KernelBody (Aliases rep)
_ Certs
_ (CoalsTab, InhibitTab)
x (PatElem (VarAliases, LetDecMem)
_, SegSpace
_, WriteReturns Certs
_ VName
return_name [(Slice SubExp, SubExp)]
_) =
case forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
return_name forall a b. (a -> b) -> a -> b
$ forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env of
Just (MemBlock PrimType
_ ShapeBase SubExp
_ VName
return_mem IxFun
_) -> (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab, InhibitTab)
x VName
return_mem
Maybe ArrayMemBound
Nothing -> forall a. HasCallStack => String -> a
error String
"Should not happen?"
makeSegMapCoals lvl -> Bool
_ lvl
_ TopdownEnv rep
td_env KernelBody (Aliases rep)
_ Certs
_ (CoalsTab, InhibitTab)
x (PatElem (VarAliases, LetDecMem)
_, SegSpace
_, KernelResult
result) =
forall a. FreeIn a => a -> Names
freeIn KernelResult
result
forall a b. a -> (a -> b) -> b
& Names -> [VName]
namesToList
forall a b. a -> (a -> b) -> b
& forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo forall a b. (a -> b) -> a -> b
$ forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env)
forall a b. a -> (a -> b) -> b
& forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall a b c. (a -> b -> c) -> b -> a -> c
flip (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal forall b c a. (b -> c) -> (a -> b) -> a -> c
. ArrayMemBound -> VName
memName) (CoalsTab, InhibitTab)
x
fullSlice :: [TPrimExp Int64 VName] -> Slice (TPrimExp Int64 VName) -> Slice (TPrimExp Int64 VName)
fullSlice :: [TPrimExp Int64 VName]
-> Slice (TPrimExp Int64 VName) -> Slice (TPrimExp Int64 VName)
fullSlice [TPrimExp Int64 VName]
shp (Slice [DimIndex (TPrimExp Int64 VName)]
slc) =
forall d. [DimIndex d] -> Slice d
Slice forall a b. (a -> b) -> a -> b
$ [DimIndex (TPrimExp Int64 VName)]
slc forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map (\TPrimExp Int64 VName
d -> forall d. d -> d -> d -> DimIndex d
DimSlice TPrimExp Int64 VName
0 TPrimExp Int64 VName
d TPrimExp Int64 VName
1) (forall a. Int -> [a] -> [a]
drop (forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimIndex (TPrimExp Int64 VName)]
slc) [TPrimExp Int64 VName]
shp)
fixPointCoalesce ::
(Coalesceable rep inner) =>
LUTabFun ->
[Param FParamMem] ->
Body (Aliases rep) ->
TopdownEnv rep ->
ShortCircuitM rep CoalsTab
fixPointCoalesce :: forall rep (inner :: * -> *).
Coalesceable rep inner =>
InhibitTab
-> [Param FParamMem]
-> Body (Aliases rep)
-> TopdownEnv rep
-> ShortCircuitM rep CoalsTab
fixPointCoalesce InhibitTab
lutab [Param FParamMem]
fpar Body (Aliases rep)
bdy TopdownEnv rep
topenv = do
BotUpEnv
buenv <- forall rep (inner :: * -> *).
Coalesceable rep inner =>
InhibitTab
-> Stms (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
mkCoalsTabStms InhibitTab
lutab (forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
bdy) TopdownEnv rep
topenv (BotUpEnv
emptyBotUpEnv {inhibit :: InhibitTab
inhibit = forall rep. TopdownEnv rep -> InhibitTab
inhibited TopdownEnv rep
topenv})
let succ_tab :: CoalsTab
succ_tab = BotUpEnv -> CoalsTab
successCoals BotUpEnv
buenv
actv_tab :: CoalsTab
actv_tab = BotUpEnv -> CoalsTab
activeCoals BotUpEnv
buenv
inhb_tab :: InhibitTab
inhb_tab = BotUpEnv -> InhibitTab
inhibit BotUpEnv
buenv
handleFunctionParams :: (CoalsTab, InhibitTab, CoalsTab)
-> (a, Uniqueness, ArrayMemBound)
-> (CoalsTab, InhibitTab, CoalsTab)
handleFunctionParams (CoalsTab
a, InhibitTab
i, CoalsTab
s) (a
_, Uniqueness
u, MemBlock PrimType
_ ShapeBase SubExp
_ VName
m IxFun
ixf) =
case (Uniqueness
u, forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m CoalsTab
a) of
(Uniqueness
Unique, Just CoalsEntry
entry)
| CoalsEntry -> IxFun
dstind CoalsEntry
entry forall a. Eq a => a -> a -> Bool
== IxFun
ixf,
Set Set LmadRef
dst_uses <- MemRefs -> AccessSummary
dstrefs (CoalsEntry -> MemRefs
memrefs CoalsEntry
entry),
Set LmadRef
dst_uses forall a. Eq a => a -> a -> Bool
== forall a. Monoid a => a
mempty ->
let (CoalsTab
a', CoalsTab
s') = (CoalsTab, CoalsTab) -> VName -> CoalsEntry -> (CoalsTab, CoalsTab)
markSuccessCoal (CoalsTab
a, CoalsTab
s) VName
m CoalsEntry
entry
in (CoalsTab
a', InhibitTab
i, CoalsTab
s')
(Uniqueness, Maybe CoalsEntry)
_ ->
let (CoalsTab
a', InhibitTab
i') = (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
a, InhibitTab
i) VName
m
in (CoalsTab
a', InhibitTab
i', CoalsTab
s)
(CoalsTab
actv_tab', InhibitTab
inhb_tab', CoalsTab
succ_tab') =
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
forall {a}.
(CoalsTab, InhibitTab, CoalsTab)
-> (a, Uniqueness, ArrayMemBound)
-> (CoalsTab, InhibitTab, CoalsTab)
handleFunctionParams
(CoalsTab
actv_tab, InhibitTab
inhb_tab, CoalsTab
succ_tab)
forall a b. (a -> b) -> a -> b
$ [Param FParamMem] -> [(VName, Uniqueness, ArrayMemBound)]
getArrMemAssocFParam [Param FParamMem]
fpar
(CoalsTab
succ_tab'', InhibitTab
failed_optdeps) = CoalsTab -> InhibitTab -> (CoalsTab, InhibitTab)
fixPointFilterDeps CoalsTab
succ_tab' forall k a. Map k a
M.empty
inhb_tab'' :: InhibitTab
inhb_tab'' = forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith forall a. Semigroup a => a -> a -> a
(<>) InhibitTab
failed_optdeps InhibitTab
inhb_tab'
if Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> Bool
M.null CoalsTab
actv_tab'
then forall a. HasCallStack => String -> a
error (String
"COALESCING ROOT: BROKEN INV, active not empty: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (forall k a. Map k a -> [k]
M.keys CoalsTab
actv_tab'))
else
if forall k a. Map k a -> Bool
M.null forall a b. (a -> b) -> a -> b
$ InhibitTab
inhb_tab'' forall k a b. Ord k => Map k a -> Map k b -> Map k a
`M.difference` forall rep. TopdownEnv rep -> InhibitTab
inhibited TopdownEnv rep
topenv
then forall (f :: * -> *) a. Applicative f => a -> f a
pure CoalsTab
succ_tab''
else forall rep (inner :: * -> *).
Coalesceable rep inner =>
InhibitTab
-> [Param FParamMem]
-> Body (Aliases rep)
-> TopdownEnv rep
-> ShortCircuitM rep CoalsTab
fixPointCoalesce InhibitTab
lutab [Param FParamMem]
fpar Body (Aliases rep)
bdy (TopdownEnv rep
topenv {inhibited :: InhibitTab
inhibited = InhibitTab
inhb_tab''})
where
fixPointFilterDeps :: CoalsTab -> InhibitTab -> (CoalsTab, InhibitTab)
fixPointFilterDeps :: CoalsTab -> InhibitTab -> (CoalsTab, InhibitTab)
fixPointFilterDeps CoalsTab
coaltab InhibitTab
inhbtab =
let (CoalsTab
coaltab', InhibitTab
inhbtab') = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
filterDeps (CoalsTab
coaltab, InhibitTab
inhbtab) (forall k a. Map k a -> [k]
M.keys CoalsTab
coaltab)
in if forall (t :: * -> *) a. Foldable t => t a -> Int
length (forall k a. Map k a -> [k]
M.keys CoalsTab
coaltab) forall a. Eq a => a -> a -> Bool
== forall (t :: * -> *) a. Foldable t => t a -> Int
length (forall k a. Map k a -> [k]
M.keys CoalsTab
coaltab')
then (CoalsTab
coaltab', InhibitTab
inhbtab')
else CoalsTab -> InhibitTab -> (CoalsTab, InhibitTab)
fixPointFilterDeps CoalsTab
coaltab' InhibitTab
inhbtab'
filterDeps :: (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
filterDeps (CoalsTab
coal, InhibitTab
inhb) VName
mb
| Bool -> Bool
not (forall k a. Ord k => k -> Map k a -> Bool
M.member VName
mb CoalsTab
coal) = (CoalsTab
coal, InhibitTab
inhb)
filterDeps (CoalsTab
coal, InhibitTab
inhb) VName
mb
| Just CoalsEntry
coal_etry <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
mb CoalsTab
coal =
let failed :: Map VName VName
failed = forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (forall {k}. Ord k => Map k CoalsEntry -> VName -> k -> Bool
failedOptDep CoalsTab
coal) (CoalsEntry -> Map VName VName
optdeps CoalsEntry
coal_etry)
in if forall k a. Map k a -> Bool
M.null Map VName VName
failed
then (CoalsTab
coal, InhibitTab
inhb)
else
(CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
coal, InhibitTab
inhb) VName
mb
filterDeps (CoalsTab, InhibitTab)
_ VName
_ = forall a. HasCallStack => String -> a
error String
"In ArrayCoalescing.hs, fun filterDeps, impossible case reached!"
failedOptDep :: Map k CoalsEntry -> VName -> k -> Bool
failedOptDep Map k CoalsEntry
coal VName
_ k
mr
| Bool -> Bool
not (k
mr forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map k CoalsEntry
coal) = Bool
True
failedOptDep Map k CoalsEntry
coal VName
r k
mr
| Just CoalsEntry
coal_etry <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
mr Map k CoalsEntry
coal = Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ VName
r forall k a. Ord k => k -> Map k a -> Bool
`M.member` CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
coal_etry
failedOptDep Map k CoalsEntry
_ VName
_ k
_ = forall a. HasCallStack => String -> a
error String
"In ArrayCoalescing.hs, fun failedOptDep, impossible case reached!"
mkCoalsTabStms ::
(Coalesceable rep inner) =>
LUTabFun ->
Stms (Aliases rep) ->
TopdownEnv rep ->
BotUpEnv ->
ShortCircuitM rep BotUpEnv
mkCoalsTabStms :: forall rep (inner :: * -> *).
Coalesceable rep inner =>
InhibitTab
-> Stms (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
mkCoalsTabStms InhibitTab
lutab Stms (Aliases rep)
stms0 = Stms (Aliases rep)
-> TopdownEnv rep -> BotUpEnv -> ShortCircuitM rep BotUpEnv
traverseStms Stms (Aliases rep)
stms0
where
non_negs_in_pats :: Names
non_negs_in_pats = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall rep. Typed rep => Pat rep -> Names
nonNegativesInPat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Stm rep -> Pat (LetDec rep)
stmPat) Stms (Aliases rep)
stms0
traverseStms :: Stms (Aliases rep)
-> TopdownEnv rep -> BotUpEnv -> ShortCircuitM rep BotUpEnv
traverseStms Stms (Aliases rep)
Empty TopdownEnv rep
_ BotUpEnv
bu_env = forall (f :: * -> *) a. Applicative f => a -> f a
pure BotUpEnv
bu_env
traverseStms (Stm (Aliases rep)
stm :<| Stms (Aliases rep)
stms) TopdownEnv rep
td_env BotUpEnv
bu_env = do
let td_env' :: TopdownEnv rep
td_env' = forall rep (inner :: * -> *).
(ASTRep rep, Op rep ~ MemOp inner rep,
TopDownHelper (inner (Aliases rep))) =>
TopdownEnv rep -> Stm (Aliases rep) -> TopdownEnv rep
updateTopdownEnv TopdownEnv rep
td_env Stm (Aliases rep)
stm
BotUpEnv
bu_env' <- Stms (Aliases rep)
-> TopdownEnv rep -> BotUpEnv -> ShortCircuitM rep BotUpEnv
traverseStms Stms (Aliases rep)
stms TopdownEnv rep
td_env' BotUpEnv
bu_env
forall rep (inner :: * -> *).
Coalesceable rep inner =>
InhibitTab
-> Stm (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
mkCoalsTabStm InhibitTab
lutab Stm (Aliases rep)
stm (TopdownEnv rep
td_env' {nonNegatives :: Names
nonNegatives = forall rep. TopdownEnv rep -> Names
nonNegatives TopdownEnv rep
td_env' forall a. Semigroup a => a -> a -> a
<> Names
non_negs_in_pats}) BotUpEnv
bu_env'
mkCoalsTabStm ::
(Coalesceable rep inner) =>
LUTabFun ->
Stm (Aliases rep) ->
TopdownEnv rep ->
BotUpEnv ->
ShortCircuitM rep BotUpEnv
mkCoalsTabStm :: forall rep (inner :: * -> *).
Coalesceable rep inner =>
InhibitTab
-> Stm (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
mkCoalsTabStm InhibitTab
_ (Let (Pat [PatElem (LetDec (Aliases rep))
pe]) StmAux (ExpDec (Aliases rep))
_ Exp (Aliases rep)
e) TopdownEnv rep
td_env BotUpEnv
bu_env
| Just PrimExp VName
primexp <- forall (m :: * -> *) rep v.
(MonadFail m, RepTypes rep) =>
(VName -> m (PrimExp v)) -> Exp rep -> m (PrimExp v)
primExpFromExp (forall rep.
AliasableRep rep =>
ScopeTab rep
-> Map VName (PrimExp VName) -> VName -> Maybe (PrimExp VName)
vnameToPrimExp (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) (BotUpEnv -> Map VName (PrimExp VName)
scals BotUpEnv
bu_env)) Exp (Aliases rep)
e =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ BotUpEnv
bu_env {scals :: Map VName (PrimExp VName)
scals = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (forall dec. PatElem dec -> VName
patElemName PatElem (LetDec (Aliases rep))
pe) PrimExp VName
primexp (BotUpEnv -> Map VName (PrimExp VName)
scals BotUpEnv
bu_env)}
mkCoalsTabStm InhibitTab
lutab (Let Pat (LetDec (Aliases rep))
patt StmAux (ExpDec (Aliases rep))
_ (Match [SubExp]
_ [Case (Body (Aliases rep))]
cases Body (Aliases rep)
defbody MatchDec (BranchType (Aliases rep))
_)) TopdownEnv rep
td_env BotUpEnv
bu_env = do
let pat_val_elms :: [PatElem (VarAliases, LetDecMem)]
pat_val_elms = forall dec. Pat dec -> [PatElem dec]
patElems Pat (LetDec (Aliases rep))
patt
(CoalsTab
activeCoals0, InhibitTab
inhibit0) =
forall rep.
HasMemBlock (Aliases rep) =>
CoalsTab
-> InhibitTab
-> Map VName (PrimExp VName)
-> TopdownEnv rep
-> [PatElem (VarAliases, LetDecMem)]
-> (CoalsTab, InhibitTab)
filterSafetyCond2and5
(BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env)
(BotUpEnv -> InhibitTab
inhibit BotUpEnv
bu_env)
(BotUpEnv -> Map VName (PrimExp VName)
scals BotUpEnv
bu_env)
TopdownEnv rep
td_env
[PatElem (VarAliases, LetDecMem)]
pat_val_elms
res_mem_def :: [MemBodyResult]
res_mem_def = forall rep.
HasMemBlock (Aliases rep) =>
CoalsTab
-> ScopeTab rep
-> [PatElem (VarAliases, LetDecMem)]
-> Body (Aliases rep)
-> [MemBodyResult]
findMemBodyResult CoalsTab
activeCoals0 (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) [PatElem (VarAliases, LetDecMem)]
pat_val_elms Body (Aliases rep)
defbody
res_mem_cases :: [[MemBodyResult]]
res_mem_cases = forall a b. (a -> b) -> [a] -> [b]
map (forall rep.
HasMemBlock (Aliases rep) =>
CoalsTab
-> ScopeTab rep
-> [PatElem (VarAliases, LetDecMem)]
-> Body (Aliases rep)
-> [MemBodyResult]
findMemBodyResult CoalsTab
activeCoals0 (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) [PatElem (VarAliases, LetDecMem)]
pat_val_elms forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall body. Case body -> body
caseBody) [Case (Body (Aliases rep))]
cases
subs_def :: FreeVarSubsts
subs_def = forall aliases.
Pat (aliases, LetDecMem) -> [SubExp] -> FreeVarSubsts
mkSubsTab Pat (LetDec (Aliases rep))
patt forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp forall a b. (a -> b) -> a -> b
$ forall rep. Body rep -> Result
bodyResult Body (Aliases rep)
defbody
subs_cases :: [FreeVarSubsts]
subs_cases = forall a b. (a -> b) -> [a] -> [b]
map (forall aliases.
Pat (aliases, LetDecMem) -> [SubExp] -> FreeVarSubsts
mkSubsTab Pat (LetDec (Aliases rep))
patt forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Body rep -> Result
bodyResult forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall body. Case body -> body
caseBody) [Case (Body (Aliases rep))]
cases
actv_def_i :: CoalsTab
actv_def_i = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (FreeVarSubsts -> CoalsTab -> MemBodyResult -> CoalsTab
transferCoalsToBody FreeVarSubsts
subs_def) CoalsTab
activeCoals0 [MemBodyResult]
res_mem_def
actv_cases_i :: [CoalsTab]
actv_cases_i = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\FreeVarSubsts
subs [MemBodyResult]
res -> forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (FreeVarSubsts -> CoalsTab -> MemBodyResult -> CoalsTab
transferCoalsToBody FreeVarSubsts
subs) CoalsTab
activeCoals0 [MemBodyResult]
res) [FreeVarSubsts]
subs_cases [[MemBodyResult]]
res_mem_cases
aux :: Map VName a -> MemBodyResult -> Map VName a
aux Map VName a
ac (MemBodyResult VName
m_b VName
_ VName
_ VName
m_r) = if VName
m_b forall a. Eq a => a -> a -> Bool
== VName
m_r then Map VName a
ac else forall k a. Ord k => k -> Map k a -> Map k a
M.delete VName
m_b Map VName a
ac
actv_def :: CoalsTab
actv_def = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall {a}. Map VName a -> MemBodyResult -> Map VName a
aux CoalsTab
actv_def_i [MemBodyResult]
res_mem_def
actv_cases :: [CoalsTab]
actv_cases = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall {a}. Map VName a -> MemBodyResult -> Map VName a
aux) [CoalsTab]
actv_cases_i [[MemBodyResult]]
res_mem_cases
BotUpEnv
res_def <- forall rep (inner :: * -> *).
Coalesceable rep inner =>
InhibitTab
-> Stms (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
mkCoalsTabStms InhibitTab
lutab (forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
defbody) TopdownEnv rep
td_env (BotUpEnv
bu_env {activeCoals :: CoalsTab
activeCoals = CoalsTab
actv_def})
[BotUpEnv]
res_cases <- forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (\Case (Body (Aliases rep))
c CoalsTab
a -> forall rep (inner :: * -> *).
Coalesceable rep inner =>
InhibitTab
-> Stms (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
mkCoalsTabStms InhibitTab
lutab (forall rep. Body rep -> Stms rep
bodyStms forall a b. (a -> b) -> a -> b
$ forall body. Case body -> body
caseBody Case (Body (Aliases rep))
c) TopdownEnv rep
td_env (BotUpEnv
bu_env {activeCoals :: CoalsTab
activeCoals = CoalsTab
a})) [Case (Body (Aliases rep))]
cases [CoalsTab]
actv_cases
let (CoalsTab
actv_def0, CoalsTab
succ_def0, InhibitTab
inhb_def0) = (BotUpEnv -> CoalsTab
activeCoals BotUpEnv
res_def, BotUpEnv -> CoalsTab
successCoals BotUpEnv
res_def, BotUpEnv -> InhibitTab
inhibit BotUpEnv
res_def)
((CoalsTab
activeCoals1, InhibitTab
inhibit1), CoalsTab
successCoals1) =
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
( forall {c}.
[(CoalsTab, Map VName c)]
-> ((CoalsTab, InhibitTab), CoalsTab)
-> [MemBodyResult]
-> ((CoalsTab, InhibitTab), CoalsTab)
foldfun
( (CoalsTab
actv_def0, CoalsTab
succ_def0)
forall a. a -> [a] -> [a]
: forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map BotUpEnv -> CoalsTab
activeCoals [BotUpEnv]
res_cases) (forall a b. (a -> b) -> [a] -> [b]
map BotUpEnv -> CoalsTab
successCoals [BotUpEnv]
res_cases)
)
)
((CoalsTab
activeCoals0, InhibitTab
inhibit0), BotUpEnv -> CoalsTab
successCoals BotUpEnv
bu_env)
(forall a. [[a]] -> [[a]]
L.transpose forall a b. (a -> b) -> a -> b
$ [MemBodyResult]
res_mem_def forall a. a -> [a] -> [a]
: [[MemBodyResult]]
res_mem_cases)
actv_res :: CoalsTab
actv_res = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith CoalsEntry -> CoalsEntry -> CoalsEntry
unionCoalsEntry) CoalsTab
activeCoals1 forall a b. (a -> b) -> a -> b
$ CoalsTab
actv_def0 forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map BotUpEnv -> CoalsTab
activeCoals [BotUpEnv]
res_cases
succ_res :: CoalsTab
succ_res = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith CoalsEntry -> CoalsEntry -> CoalsEntry
unionCoalsEntry) CoalsTab
successCoals1 forall a b. (a -> b) -> a -> b
$ CoalsTab
succ_def0 forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map BotUpEnv -> CoalsTab
successCoals [BotUpEnv]
res_cases
inhibit_res :: InhibitTab
inhibit_res =
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
M.unionsWith
forall a. Semigroup a => a -> a -> a
(<>)
( InhibitTab
inhibit1
forall a. a -> [a] -> [a]
: forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
( \CoalsTab
actv InhibitTab
inhb ->
let failed :: CoalsTab
failed = forall k a b. Ord k => Map k a -> Map k b -> Map k a
M.difference CoalsTab
actv forall a b. (a -> b) -> a -> b
$ forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith CoalsEntry -> CoalsEntry -> CoalsEntry
unionCoalsEntry CoalsTab
actv CoalsTab
activeCoals0
in forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
failed, InhibitTab
inhb) (forall k a. Map k a -> [k]
M.keys CoalsTab
failed)
)
(CoalsTab
actv_def0 forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map BotUpEnv -> CoalsTab
activeCoals [BotUpEnv]
res_cases)
(InhibitTab
inhb_def0 forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map BotUpEnv -> InhibitTab
inhibit [BotUpEnv]
res_cases)
)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
BotUpEnv
bu_env
{ activeCoals :: CoalsTab
activeCoals = CoalsTab
actv_res,
successCoals :: CoalsTab
successCoals = CoalsTab
succ_res,
inhibit :: InhibitTab
inhibit = InhibitTab
inhibit_res
}
where
foldfun :: [(CoalsTab, Map VName c)]
-> ((CoalsTab, InhibitTab), CoalsTab)
-> [MemBodyResult]
-> ((CoalsTab, InhibitTab), CoalsTab)
foldfun [(CoalsTab, Map VName c)]
_ ((CoalsTab, InhibitTab), CoalsTab)
_ [] =
forall a. HasCallStack => String -> a
error String
"Imposible Case 1!!!"
foldfun [(CoalsTab, Map VName c)]
_ ((CoalsTab
act, InhibitTab
_), CoalsTab
_) [MemBodyResult]
mem_body_results
| Maybe CoalsEntry
Nothing <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (MemBodyResult -> VName
patMem forall a b. (a -> b) -> a -> b
$ forall a. [a] -> a
head [MemBodyResult]
mem_body_results) CoalsTab
act =
forall a. HasCallStack => String -> a
error String
"Imposible Case 2!!!"
foldfun
[(CoalsTab, Map VName c)]
acc
((CoalsTab
act, InhibitTab
inhb), CoalsTab
succc)
mem_body_results :: [MemBodyResult]
mem_body_results@(MemBodyResult VName
m_b VName
_ VName
_ VName
_ : [MemBodyResult]
_)
| Just CoalsEntry
info <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_b CoalsTab
act,
Just [c]
_ <- forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup forall b c a. (b -> c) -> (a -> b) -> a -> c
. MemBodyResult -> VName
bodyMem) [MemBodyResult]
mem_body_results forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(CoalsTab, Map VName c)]
acc =
let info' :: CoalsEntry
info' =
CoalsEntry
info
{ optdeps :: Map VName VName
optdeps =
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
(\MemBodyResult
mbr -> forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (MemBodyResult -> VName
bodyName MemBodyResult
mbr) (MemBodyResult -> VName
bodyMem MemBodyResult
mbr))
(CoalsEntry -> Map VName VName
optdeps CoalsEntry
info)
[MemBodyResult]
mem_body_results
}
(CoalsTab
act', CoalsTab
succc') = (CoalsTab, CoalsTab) -> VName -> CoalsEntry -> (CoalsTab, CoalsTab)
markSuccessCoal (CoalsTab
act, CoalsTab
succc) VName
m_b CoalsEntry
info'
in ((CoalsTab
act', InhibitTab
inhb), CoalsTab
succc')
foldfun
[(CoalsTab, Map VName c)]
acc
((CoalsTab
act, InhibitTab
inhb), CoalsTab
succc)
mem_body_results :: [MemBodyResult]
mem_body_results@(MemBodyResult VName
m_b VName
_ VName
_ VName
_ : [MemBodyResult]
_)
| Just CoalsEntry
info <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_b CoalsTab
act,
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall a. Eq a => a -> a -> Bool
(==) VName
m_b forall b c a. (b -> c) -> (a -> b) -> a -> c
. MemBodyResult -> VName
bodyMem) [MemBodyResult]
mem_body_results,
Just [CoalsEntry]
info' <- forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup forall b c a. (b -> c) -> (a -> b) -> a -> c
. MemBodyResult -> VName
bodyMem) [MemBodyResult]
mem_body_results forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(CoalsTab, Map VName c)]
acc =
let info'' :: CoalsEntry
info'' =
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl CoalsEntry -> CoalsEntry -> CoalsEntry
unionCoalsEntry CoalsEntry
info [CoalsEntry]
info'
act' :: CoalsTab
act' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
m_b CoalsEntry
info'' CoalsTab
act
in ((CoalsTab
act', InhibitTab
inhb), CoalsTab
succc)
foldfun [(CoalsTab, Map VName c)]
_ ((CoalsTab
act, InhibitTab
inhb), CoalsTab
succc) (MemBodyResult
mbr : [MemBodyResult]
_) =
((CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
act, InhibitTab
inhb) (MemBodyResult -> VName
patMem MemBodyResult
mbr), CoalsTab
succc)
mkCoalsTabStm InhibitTab
lutab (Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
_ (Loop [(FParam (Aliases rep), SubExp)]
arginis LoopForm
lform Body (Aliases rep)
body)) TopdownEnv rep
td_env BotUpEnv
bu_env = do
let pat_val_elms :: [PatElem (VarAliases, LetDecMem)]
pat_val_elms = forall dec. Pat dec -> [PatElem dec]
patElems Pat (LetDec (Aliases rep))
pat
(CoalsTab
actv0, InhibitTab
inhibit0) =
forall rep.
HasMemBlock (Aliases rep) =>
CoalsTab
-> InhibitTab
-> Map VName (PrimExp VName)
-> TopdownEnv rep
-> [PatElem (VarAliases, LetDecMem)]
-> (CoalsTab, InhibitTab)
filterSafetyCond2and5
(BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env)
(BotUpEnv -> InhibitTab
inhibit BotUpEnv
bu_env)
(BotUpEnv -> Map VName (PrimExp VName)
scals BotUpEnv
bu_env)
TopdownEnv rep
td_env
[PatElem (VarAliases, LetDecMem)]
pat_val_elms
bdy_ress :: Result
bdy_ress = forall rep. Body rep -> Result
bodyResult Body (Aliases rep)
body
([(VName, VName)]
patmems, [(VName, VName)]
argmems, [(VName, VName)]
inimems, [(VName, VName)]
resmems) =
forall a b c d. [(a, b, c, d)] -> ([a], [b], [c], [d])
L.unzip4 forall a b. (a -> b) -> a -> b
$
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (CoalsTab
-> (PatElem (VarAliases, LetDecMem), (Param FParamMem, SubExp),
SubExp)
-> Maybe
((VName, VName), (VName, VName), (VName, VName), (VName, VName))
mapmbFun CoalsTab
actv0) (forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [PatElem (VarAliases, LetDecMem)]
pat_val_elms [(FParam (Aliases rep), SubExp)]
arginis forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp Result
bdy_ress)
coal_pat_names :: Names
coal_pat_names = [VName] -> Names
namesFromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(VName, VName)]
patmems
(CoalsTab
actv1, InhibitTab
inhibit1) =
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
( \(CoalsTab
act, InhibitTab
inhb) (VName
b, MemBlock PrimType
_ ShapeBase SubExp
_ VName
m_b IxFun
_) ->
if VName
b VName -> Names -> Bool
`nameIn` Names
coal_pat_names
then (CoalsTab
act, InhibitTab
inhb)
else (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
act, InhibitTab
inhb) VName
m_b
)
(CoalsTab
actv0, InhibitTab
inhibit0)
(forall aliases.
Pat (aliases, LetDecMem) -> [(VName, ArrayMemBound)]
getArrMemAssoc Pat (LetDec (Aliases rep))
pat)
res_mem_bdy :: [MemBodyResult]
res_mem_bdy = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(VName
b, VName
m_b) (VName
r, VName
m_r) -> VName -> VName -> VName -> VName -> MemBodyResult
MemBodyResult VName
m_b VName
b VName
r VName
m_r) [(VName, VName)]
patmems [(VName, VName)]
resmems
res_mem_arg :: [MemBodyResult]
res_mem_arg = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(VName
b, VName
m_b) (VName
r, VName
m_r) -> VName -> VName -> VName -> VName -> MemBodyResult
MemBodyResult VName
m_b VName
b VName
r VName
m_r) [(VName, VName)]
patmems [(VName, VName)]
argmems
res_mem_ini :: [MemBodyResult]
res_mem_ini = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(VName
b, VName
m_b) (VName
r, VName
m_r) -> VName -> VName -> VName -> VName -> MemBodyResult
MemBodyResult VName
m_b VName
b VName
r VName
m_r) [(VName, VName)]
patmems [(VName, VName)]
inimems
actv2 :: CoalsTab
actv2 =
let subs_res :: FreeVarSubsts
subs_res = forall aliases.
Pat (aliases, LetDecMem) -> [SubExp] -> FreeVarSubsts
mkSubsTab Pat (LetDec (Aliases rep))
pat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp forall a b. (a -> b) -> a -> b
$ forall rep. Body rep -> Result
bodyResult Body (Aliases rep)
body
actv11 :: CoalsTab
actv11 = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (FreeVarSubsts -> CoalsTab -> MemBodyResult -> CoalsTab
transferCoalsToBody FreeVarSubsts
subs_res) CoalsTab
actv1 [MemBodyResult]
res_mem_bdy
subs_arg :: FreeVarSubsts
subs_arg = forall aliases.
Pat (aliases, LetDecMem) -> [SubExp] -> FreeVarSubsts
mkSubsTab Pat (LetDec (Aliases rep))
pat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
Var forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. Param dec -> VName
paramName forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(FParam (Aliases rep), SubExp)]
arginis
actv12 :: CoalsTab
actv12 = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (FreeVarSubsts -> CoalsTab -> MemBodyResult -> CoalsTab
transferCoalsToBody FreeVarSubsts
subs_arg) CoalsTab
actv11 [MemBodyResult]
res_mem_arg
subs_ini :: FreeVarSubsts
subs_ini = forall aliases.
Pat (aliases, LetDecMem) -> [SubExp] -> FreeVarSubsts
mkSubsTab Pat (LetDec (Aliases rep))
pat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(FParam (Aliases rep), SubExp)]
arginis
in forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (FreeVarSubsts -> CoalsTab -> MemBodyResult -> CoalsTab
transferCoalsToBody FreeVarSubsts
subs_ini) CoalsTab
actv12 [MemBodyResult]
res_mem_ini
insertMemAliases :: CoalsTab -> (MemBodyResult, MemBodyResult) -> CoalsTab
insertMemAliases CoalsTab
tab (MemBodyResult VName
_ VName
_ VName
_ VName
m_r, MemBodyResult VName
_ VName
_ VName
_ VName
m_a) =
if VName
m_r forall a. Eq a => a -> a -> Bool
== VName
m_a
then CoalsTab
tab
else case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_r CoalsTab
tab of
Maybe CoalsEntry
Nothing -> CoalsTab
tab
Just CoalsEntry
etry ->
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
m_r (CoalsEntry
etry {alsmem :: Names
alsmem = CoalsEntry -> Names
alsmem CoalsEntry
etry forall a. Semigroup a => a -> a -> a
<> VName -> Names
oneName VName
m_a}) CoalsTab
tab
actv3 :: CoalsTab
actv3 = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl CoalsTab -> (MemBodyResult, MemBodyResult) -> CoalsTab
insertMemAliases CoalsTab
actv2 (forall a b. [a] -> [b] -> [(a, b)]
zip [MemBodyResult]
res_mem_bdy [MemBodyResult]
res_mem_arg)
actv4 :: CoalsTab
actv4 = forall a b k. (a -> b) -> Map k a -> Map k b
M.map (\CoalsEntry
etry -> CoalsEntry
etry {memrefs :: MemRefs
memrefs = forall a. Monoid a => a
mempty}) CoalsTab
actv3
BotUpEnv
res_env_body <-
forall rep (inner :: * -> *).
Coalesceable rep inner =>
InhibitTab
-> Stms (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
mkCoalsTabStms
InhibitTab
lutab
(forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
body)
TopdownEnv rep
td_env'
( BotUpEnv
bu_env
{ activeCoals :: CoalsTab
activeCoals = CoalsTab
actv4,
inhibit :: InhibitTab
inhibit = InhibitTab
inhibit1
}
)
let scals_loop :: Map VName (PrimExp VName)
scals_loop = BotUpEnv -> Map VName (PrimExp VName)
scals BotUpEnv
res_env_body
(CoalsTab
res_actv0, CoalsTab
res_succ0, InhibitTab
res_inhb0) = (BotUpEnv -> CoalsTab
activeCoals BotUpEnv
res_env_body, BotUpEnv -> CoalsTab
successCoals BotUpEnv
res_env_body, BotUpEnv -> InhibitTab
inhibit BotUpEnv
res_env_body)
mb_loop_idx :: Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
mb_loop_idx = LoopForm
-> Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
mbLoopIndexRange LoopForm
lform
CoalsTab
res_actv1 <- forall k (m :: * -> *) v.
(Eq k, Monad m) =>
(v -> m Bool) -> Map k v -> m (Map k v)
filterMapM1 (Map VName (PrimExp VName)
-> Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
-> CoalsEntry
-> ShortCircuitM rep Bool
loopSoundness1Entry Map VName (PrimExp VName)
scals_loop Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
mb_loop_idx) CoalsTab
res_actv0
CoalsTab
res_actv2 <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (ScopeTab rep
-> Map VName (PrimExp VName)
-> Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
-> CoalsEntry
-> ShortCircuitM rep CoalsEntry
aggAcrossLoopEntry (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env' forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf (forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
body)) Map VName (PrimExp VName)
scals_loop Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
mb_loop_idx) CoalsTab
res_actv1
let res_actv3 :: CoalsTab
res_actv3 = forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (CoalsTab -> VName -> CoalsEntry -> Bool
loopSoundness2Entry CoalsTab
actv3) CoalsTab
res_actv2
let tmp_succ :: CoalsTab
tmp_succ =
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (forall {k} {a} {p}. Ord k => Map k a -> k -> p -> Bool
okLookup CoalsTab
actv3) forall a b. (a -> b) -> a -> b
$
forall k a b. Ord k => Map k a -> Map k b -> Map k a
M.difference CoalsTab
res_succ0 (BotUpEnv -> CoalsTab
successCoals BotUpEnv
bu_env)
ver_succ :: CoalsTab
ver_succ = forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (CoalsTab -> VName -> CoalsEntry -> Bool
loopSoundness2Entry CoalsTab
actv3) CoalsTab
tmp_succ
let suc_fail :: CoalsTab
suc_fail = forall k a b. Ord k => Map k a -> Map k b -> Map k a
M.difference CoalsTab
tmp_succ CoalsTab
ver_succ
(CoalsTab
res_succ, InhibitTab
res_inhb1) = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
res_succ0, InhibitTab
res_inhb0) forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [k]
M.keys CoalsTab
suc_fail
act_fail :: CoalsTab
act_fail = forall k a b. Ord k => Map k a -> Map k b -> Map k a
M.difference CoalsTab
res_actv0 CoalsTab
res_actv3
(CoalsTab
_, InhibitTab
res_inhb) = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
res_actv0, InhibitTab
res_inhb1) forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [k]
M.keys CoalsTab
act_fail
res_actv :: CoalsTab
res_actv =
forall k a b. (k -> a -> b) -> Map k a -> Map k b
M.mapWithKey (forall {k}.
Ord k =>
Map k CoalsEntry -> k -> CoalsEntry -> CoalsEntry
addBeforeLoop CoalsTab
actv3) CoalsTab
res_actv3
((CoalsTab
fin_actv1, InhibitTab
fin_inhb1), CoalsTab
fin_succ1) =
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ((CoalsTab, InhibitTab), CoalsTab)
-> ((VName, VName), (VName, VName), (VName, VName), (VName, VName))
-> ((CoalsTab, InhibitTab), CoalsTab)
foldFunOptimPromotion ((CoalsTab
res_actv, InhibitTab
res_inhb), CoalsTab
res_succ) forall a b. (a -> b) -> a -> b
$
forall a b c d. [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]
L.zip4 [(VName, VName)]
patmems [(VName, VName)]
argmems [(VName, VName)]
resmems [(VName, VName)]
inimems
(CoalsTab
fin_actv2, InhibitTab
fin_inhb2) =
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey
( \(CoalsTab, InhibitTab)
acc VName
k CoalsEntry
_ ->
if VName
k VName -> Names -> Bool
`nameIn` [VName] -> Names
namesFromList (forall a b. (a -> b) -> [a] -> [b]
map (forall dec. Param dec -> VName
paramName forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(FParam (Aliases rep), SubExp)]
arginis)
then (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab, InhibitTab)
acc VName
k
else (CoalsTab, InhibitTab)
acc
)
(CoalsTab
fin_actv1, InhibitTab
fin_inhb1)
CoalsTab
fin_actv1
forall (f :: * -> *) a. Applicative f => a -> f a
pure BotUpEnv
bu_env {activeCoals :: CoalsTab
activeCoals = CoalsTab
fin_actv2, successCoals :: CoalsTab
successCoals = CoalsTab
fin_succ1, inhibit :: InhibitTab
inhibit = InhibitTab
fin_inhb2}
where
allocs_bdy :: AllocTab
allocs_bdy = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall {rep} {inner :: * -> *}.
(OpC rep ~ MemOp inner) =>
AllocTab -> Stm rep -> AllocTab
getAllocs (forall rep. TopdownEnv rep -> AllocTab
alloc TopdownEnv rep
td_env') forall a b. (a -> b) -> a -> b
$ forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
body
td_env_allocs :: TopdownEnv rep
td_env_allocs = TopdownEnv rep
td_env' {alloc :: AllocTab
alloc = AllocTab
allocs_bdy, scope :: ScopeTab rep
scope = forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env' forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf (forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
body)}
td_env' :: TopdownEnv rep
td_env' = forall rep.
TopdownEnv rep
-> [(FParam rep, SubExp)] -> LoopForm -> TopdownEnv rep
updateTopdownEnvLoop TopdownEnv rep
td_env [(FParam (Aliases rep), SubExp)]
arginis LoopForm
lform
getAllocs :: AllocTab -> Stm rep -> AllocTab
getAllocs AllocTab
tab (Let (Pat [PatElem (LetDec rep)
pe]) StmAux (ExpDec rep)
_ (Op (Alloc SubExp
_ Space
sp))) =
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (forall dec. PatElem dec -> VName
patElemName PatElem (LetDec rep)
pe) Space
sp AllocTab
tab
getAllocs AllocTab
tab Stm rep
_ = AllocTab
tab
okLookup :: Map k a -> k -> p -> Bool
okLookup Map k a
tab k
m p
_
| Just a
_ <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
m Map k a
tab = Bool
True
okLookup Map k a
_ k
_ p
_ = Bool
False
mapmbFun :: CoalsTab
-> (PatElem (VarAliases, LetDecMem), (Param FParamMem, SubExp),
SubExp)
-> Maybe
((VName, VName), (VName, VName), (VName, VName), (VName, VName))
mapmbFun CoalsTab
actv0 (PatElem (VarAliases, LetDecMem)
patel, (Param FParamMem
arg, SubExp
ini), SubExp
bdyres)
| VName
b <- forall dec. PatElem dec -> VName
patElemName PatElem (VarAliases, LetDecMem)
patel,
(VarAliases
_, MemArray PrimType
_ ShapeBase SubExp
_ NoUniqueness
_ (ArrayIn VName
m_b IxFun
_)) <- forall dec. PatElem dec -> dec
patElemDec PatElem (VarAliases, LetDecMem)
patel,
VName
a <- forall dec. Param dec -> VName
paramName Param FParamMem
arg,
Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((VName -> Names -> Bool
`nameIn` forall a. FreeIn a => a -> Names
freeIn (forall dec. Param dec -> dec
paramDec Param FParamMem
arg)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. Param dec -> VName
paramName forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(FParam (Aliases rep), SubExp)]
arginis,
Var VName
a0 <- SubExp
ini,
Var VName
r <- SubExp
bdyres,
Just CoalsEntry
coal_etry <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_b CoalsTab
actv0,
Just Coalesced
_ <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
b (CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
coal_etry),
Just (MemBlock PrimType
_ ShapeBase SubExp
_ VName
m_a IxFun
_) <- forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
a (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env_allocs),
Just (MemBlock PrimType
_ ShapeBase SubExp
_ VName
m_a0 IxFun
_) <- forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
a0 (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env_allocs),
Just (MemBlock PrimType
_ ShapeBase SubExp
_ VName
m_r IxFun
_) <- forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
r (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env_allocs),
Just Names
nms <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
a InhibitTab
lutab,
VName
a0 VName -> Names -> Bool
`nameIn` Names
nms,
VName
m_r forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` forall k a. Map k a -> [k]
M.keys (forall rep. TopdownEnv rep -> AllocTab
alloc TopdownEnv rep
td_env_allocs) =
forall a. a -> Maybe a
Just ((VName
b, VName
m_b), (VName
a, VName
m_a), (VName
a0, VName
m_a0), (VName
r, VName
m_r))
mapmbFun CoalsTab
_ (PatElem (VarAliases, LetDecMem)
_patel, (Param FParamMem
_arg, SubExp
_ini), SubExp
_bdyres) = forall a. Maybe a
Nothing
foldFunOptimPromotion ::
((CoalsTab, InhibitTab), CoalsTab) ->
((VName, VName), (VName, VName), (VName, VName), (VName, VName)) ->
((CoalsTab, InhibitTab), CoalsTab)
foldFunOptimPromotion :: ((CoalsTab, InhibitTab), CoalsTab)
-> ((VName, VName), (VName, VName), (VName, VName), (VName, VName))
-> ((CoalsTab, InhibitTab), CoalsTab)
foldFunOptimPromotion ((CoalsTab
act, InhibitTab
inhb), CoalsTab
succc) ((VName
b, VName
m_b), (VName
a, VName
m_a), (VName
_r, VName
m_r), (VName
b_i, VName
m_i))
| VName
m_r forall a. Eq a => a -> a -> Bool
== VName
m_i,
Just CoalsEntry
info <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_i CoalsTab
act,
Just Map VName Coalesced
vtab_i <- forall rep.
HasMemBlock (Aliases rep) =>
TopdownEnv rep
-> Map VName Coalesced -> VName -> Maybe (Map VName Coalesced)
addInvAliasesVarTab TopdownEnv rep
td_env (CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
info) VName
b_i =
forall a. HasCallStack => Bool -> a -> a
Exc.assert
(VName
m_r forall a. Eq a => a -> a -> Bool
== VName
m_b Bool -> Bool -> Bool
&& VName
m_a forall a. Eq a => a -> a -> Bool
== VName
m_b)
((forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
m_b (CoalsEntry
info {vartab :: Map VName Coalesced
vartab = Map VName Coalesced
vtab_i}) CoalsTab
act, InhibitTab
inhb), CoalsTab
succc)
| VName
m_r forall a. Eq a => a -> a -> Bool
== VName
m_i =
forall a. HasCallStack => Bool -> a -> a
Exc.assert
(VName
m_r forall a. Eq a => a -> a -> Bool
== VName
m_b Bool -> Bool -> Bool
&& VName
m_a forall a. Eq a => a -> a -> Bool
== VName
m_b)
((CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
act, InhibitTab
inhb) VName
m_b, CoalsTab
succc)
| Just CoalsEntry
info_b0 <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_b CoalsTab
act,
Just CoalsEntry
info_a0 <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_a CoalsTab
act,
Just CoalsEntry
info_i <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_i CoalsTab
act,
forall k a. Ord k => k -> Map k a -> Bool
M.member VName
m_r CoalsTab
succc,
Just Map VName Coalesced
vtab_i <- forall rep.
HasMemBlock (Aliases rep) =>
TopdownEnv rep
-> Map VName Coalesced -> VName -> Maybe (Map VName Coalesced)
addInvAliasesVarTab TopdownEnv rep
td_env (CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
info_i) VName
b_i,
[Just CoalsEntry
info_b, Just CoalsEntry
info_a] <- forall a b. (a -> b) -> [a] -> [b]
map (VName, CoalsEntry) -> Maybe CoalsEntry
translateIxFnInScope [(VName
b, CoalsEntry
info_b0), (VName
a, CoalsEntry
info_a0)] =
let info_b' :: CoalsEntry
info_b' = CoalsEntry
info_b {optdeps :: Map VName VName
optdeps = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
b_i VName
m_i forall a b. (a -> b) -> a -> b
$ CoalsEntry -> Map VName VName
optdeps CoalsEntry
info_b}
info_a' :: CoalsEntry
info_a' = CoalsEntry
info_a {optdeps :: Map VName VName
optdeps = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
b_i VName
m_i forall a b. (a -> b) -> a -> b
$ CoalsEntry -> Map VName VName
optdeps CoalsEntry
info_a}
info_i' :: CoalsEntry
info_i' =
CoalsEntry
info_i
{ optdeps :: Map VName VName
optdeps = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
b VName
m_b forall a b. (a -> b) -> a -> b
$ CoalsEntry -> Map VName VName
optdeps CoalsEntry
info_i,
memrefs :: MemRefs
memrefs = forall a. Monoid a => a
mempty,
vartab :: Map VName Coalesced
vartab = Map VName Coalesced
vtab_i
}
act' :: CoalsTab
act' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
m_i CoalsEntry
info_i' CoalsTab
act
(CoalsTab
act1, CoalsTab
succc1) =
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
(\(CoalsTab, CoalsTab)
acc (VName
m, CoalsEntry
info) -> (CoalsTab, CoalsTab) -> VName -> CoalsEntry -> (CoalsTab, CoalsTab)
markSuccessCoal (CoalsTab, CoalsTab)
acc VName
m CoalsEntry
info)
(CoalsTab
act', CoalsTab
succc)
[(VName
m_b, CoalsEntry
info_b'), (VName
m_a, CoalsEntry
info_a')]
in
((CoalsTab
act1, InhibitTab
inhb), CoalsTab
succc1)
foldFunOptimPromotion ((CoalsTab
act, InhibitTab
inhb), CoalsTab
succc) ((VName
_, VName
m_b), (VName
_a, VName
m_a), (VName
_r, VName
m_r), (VName
_b_i, VName
m_i)) =
forall a. HasCallStack => Bool -> a -> a
Exc.assert
(VName
m_r forall a. Eq a => a -> a -> Bool
/= VName
m_i)
(forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
act, InhibitTab
inhb) [VName
m_b, VName
m_a, VName
m_r, VName
m_i], CoalsTab
succc)
translateIxFnInScope :: (VName, CoalsEntry) -> Maybe CoalsEntry
translateIxFnInScope (VName
x, CoalsEntry
info)
| Just (Coalesced CoalescedKind
knd mbd :: ArrayMemBound
mbd@(MemBlock PrimType
_ ShapeBase SubExp
_ VName
_ IxFun
ixfn) FreeVarSubsts
_subs0) <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x (CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
info),
forall rep. TopdownEnv rep -> VName -> Bool
isInScope TopdownEnv rep
td_env (CoalsEntry -> VName
dstmem CoalsEntry
info) =
let scope_tab :: ScopeTab rep
scope_tab =
forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env
forall a. Semigroup a => a -> a -> a
<> forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(FParam (Aliases rep), SubExp)]
arginis)
in case forall a rep.
FreeIn a =>
ScopeTab rep
-> Map VName (PrimExp VName) -> a -> Maybe FreeVarSubsts
freeVarSubstitutions ScopeTab rep
scope_tab (BotUpEnv -> Map VName (PrimExp VName)
scals BotUpEnv
bu_env) IxFun
ixfn of
Just FreeVarSubsts
fv_subst ->
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ CoalsEntry
info {vartab :: Map VName Coalesced
vartab = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
x (CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced CoalescedKind
knd ArrayMemBound
mbd FreeVarSubsts
fv_subst) (CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
info)}
Maybe FreeVarSubsts
Nothing -> forall a. Maybe a
Nothing
translateIxFnInScope (VName, CoalsEntry)
_ = forall a. Maybe a
Nothing
se0 :: SubExp
se0 = IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
mbLoopIndexRange ::
LoopForm ->
Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
mbLoopIndexRange :: LoopForm
-> Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
mbLoopIndexRange (WhileLoop VName
_) = forall a. Maybe a
Nothing
mbLoopIndexRange (ForLoop VName
inm IntType
_inttp SubExp
seN) = forall a. a -> Maybe a
Just (VName
inm, (SubExp -> TPrimExp Int64 VName
pe64 SubExp
se0, SubExp -> TPrimExp Int64 VName
pe64 SubExp
seN))
addBeforeLoop :: Map k CoalsEntry -> k -> CoalsEntry -> CoalsEntry
addBeforeLoop Map k CoalsEntry
actv_bef k
m_b CoalsEntry
etry =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
m_b Map k CoalsEntry
actv_bef of
Maybe CoalsEntry
Nothing -> CoalsEntry
etry
Just CoalsEntry
etry0 ->
CoalsEntry
etry {memrefs :: MemRefs
memrefs = CoalsEntry -> MemRefs
memrefs CoalsEntry
etry0 forall a. Semigroup a => a -> a -> a
<> CoalsEntry -> MemRefs
memrefs CoalsEntry
etry}
aggAcrossLoopEntry :: ScopeTab rep
-> Map VName (PrimExp VName)
-> Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
-> CoalsEntry
-> ShortCircuitM rep CoalsEntry
aggAcrossLoopEntry ScopeTab rep
scope_loop Map VName (PrimExp VName)
scal_tab Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
idx CoalsEntry
etry = do
AccessSummary
wrts <-
forall (m :: * -> *) rep.
MonadFreshNames m =>
ScopeTab rep
-> ScopeTab rep
-> Map VName (PrimExp VName)
-> Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
-> AccessSummary
-> m AccessSummary
aggSummaryLoopTotal (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) ScopeTab rep
scope_loop Map VName (PrimExp VName)
scal_tab Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
idx forall a b. (a -> b) -> a -> b
$
(MemRefs -> AccessSummary
srcwrts forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoalsEntry -> MemRefs
memrefs) CoalsEntry
etry
AccessSummary
uses <-
forall (m :: * -> *) rep.
MonadFreshNames m =>
ScopeTab rep
-> ScopeTab rep
-> Map VName (PrimExp VName)
-> Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
-> AccessSummary
-> m AccessSummary
aggSummaryLoopTotal (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) ScopeTab rep
scope_loop Map VName (PrimExp VName)
scal_tab Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
idx forall a b. (a -> b) -> a -> b
$
(MemRefs -> AccessSummary
dstrefs forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoalsEntry -> MemRefs
memrefs) CoalsEntry
etry
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ CoalsEntry
etry {memrefs :: MemRefs
memrefs = AccessSummary -> AccessSummary -> MemRefs
MemRefs AccessSummary
uses AccessSummary
wrts}
loopSoundness1Entry :: Map VName (PrimExp VName)
-> Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
-> CoalsEntry
-> ShortCircuitM rep Bool
loopSoundness1Entry Map VName (PrimExp VName)
scal_tab Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
idx CoalsEntry
etry = do
let wrt_i :: AccessSummary
wrt_i = (MemRefs -> AccessSummary
srcwrts forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoalsEntry -> MemRefs
memrefs) CoalsEntry
etry
AccessSummary
use_p <-
forall (m :: * -> *).
MonadFreshNames m =>
Map VName (PrimExp VName)
-> Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
-> AccessSummary
-> m AccessSummary
aggSummaryLoopPartial (Map VName (PrimExp VName)
scal_tab forall a. Semigroup a => a -> a -> a
<> forall rep. TopdownEnv rep -> Map VName (PrimExp VName)
scalarTable TopdownEnv rep
td_env) Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
idx forall a b. (a -> b) -> a -> b
$
MemRefs -> AccessSummary
dstrefs forall a b. (a -> b) -> a -> b
$
CoalsEntry -> MemRefs
memrefs CoalsEntry
etry
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall rep.
AliasableRep rep =>
TopdownEnv rep -> AccessSummary -> AccessSummary -> Bool
noMemOverlap TopdownEnv rep
td_env' AccessSummary
wrt_i AccessSummary
use_p
loopSoundness2Entry :: CoalsTab -> VName -> CoalsEntry -> Bool
loopSoundness2Entry :: CoalsTab -> VName -> CoalsEntry -> Bool
loopSoundness2Entry CoalsTab
old_actv VName
m_b CoalsEntry
etry =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_b CoalsTab
old_actv of
Maybe CoalsEntry
Nothing -> Bool
True
Just CoalsEntry
etry0 ->
let uses_before :: AccessSummary
uses_before = (MemRefs -> AccessSummary
dstrefs forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoalsEntry -> MemRefs
memrefs) CoalsEntry
etry0
write_loop :: AccessSummary
write_loop = (MemRefs -> AccessSummary
srcwrts forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoalsEntry -> MemRefs
memrefs) CoalsEntry
etry
in forall rep.
AliasableRep rep =>
TopdownEnv rep -> AccessSummary -> AccessSummary -> Bool
noMemOverlap TopdownEnv rep
td_env AccessSummary
write_loop AccessSummary
uses_before
mkCoalsTabStm InhibitTab
lutab stm :: Stm (Aliases rep)
stm@(Let pat :: Pat (LetDec (Aliases rep))
pat@(Pat [PatElem (LetDec (Aliases rep))
x']) StmAux (ExpDec (Aliases rep))
_ (BasicOp (Update Safety
safety VName
x Slice SubExp
_ SubExp
_elm))) TopdownEnv rep
td_env BotUpEnv
bu_env
| [(VName
_, MemBlock PrimType
_ ShapeBase SubExp
_ VName
m_x IxFun
_)] <- forall aliases.
Pat (aliases, LetDecMem) -> [(VName, ArrayMemBound)]
getArrMemAssoc Pat (LetDec (Aliases rep))
pat = do
let (CoalsTab
actv, InhibitTab
inhbt) = forall rep (inner :: * -> *).
(AliasableRep rep, Op rep ~ MemOp inner rep,
HasMemBlock (Aliases rep)) =>
TopdownEnv rep
-> BotUpEnv -> Stm (Aliases rep) -> (CoalsTab, InhibitTab)
recordMemRefUses TopdownEnv rep
td_env BotUpEnv
bu_env Stm (Aliases rep)
stm
(CoalsTab
actv', InhibitTab
inhbt') =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_x CoalsTab
actv of
Maybe CoalsEntry
Nothing -> (CoalsTab
actv, InhibitTab
inhbt)
Just CoalsEntry
info ->
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall dec. PatElem dec -> VName
patElemName PatElem (LetDec (Aliases rep))
x') (CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
info) of
Maybe Coalesced
Nothing -> (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
actv, InhibitTab
inhbt) VName
m_x
Just (Coalesced CoalescedKind
k mblk :: ArrayMemBound
mblk@(MemBlock PrimType
_ ShapeBase SubExp
_ VName
_ IxFun
x_indfun) FreeVarSubsts
_) ->
case forall a rep.
FreeIn a =>
ScopeTab rep
-> Map VName (PrimExp VName) -> a -> Maybe FreeVarSubsts
freeVarSubstitutions (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) (BotUpEnv -> Map VName (PrimExp VName)
scals BotUpEnv
bu_env) IxFun
x_indfun of
Just FreeVarSubsts
fv_subs
| forall rep. TopdownEnv rep -> VName -> Bool
isInScope TopdownEnv rep
td_env (CoalsEntry -> VName
dstmem CoalsEntry
info) ->
let coal_etry_x :: Coalesced
coal_etry_x = CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced CoalescedKind
k ArrayMemBound
mblk FreeVarSubsts
fv_subs
info' :: CoalsEntry
info' =
CoalsEntry
info
{ vartab :: Map VName Coalesced
vartab =
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
x Coalesced
coal_etry_x forall a b. (a -> b) -> a -> b
$
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (forall dec. PatElem dec -> VName
patElemName PatElem (LetDec (Aliases rep))
x') Coalesced
coal_etry_x (CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
info)
}
in (forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
m_x CoalsEntry
info' CoalsTab
actv, InhibitTab
inhbt)
Maybe FreeVarSubsts
_ ->
(CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
actv, InhibitTab
inhbt) VName
m_x
CoalsTab
actv'' <-
if Safety
safety forall a. Eq a => a -> a -> Bool
== Safety
Unsafe
then forall rep (inner :: * -> *).
Coalesceable rep inner =>
Stm (Aliases rep)
-> InhibitTab
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep CoalsTab
mkCoalsHelper3PatternMatch Stm (Aliases rep)
stm InhibitTab
lutab TopdownEnv rep
td_env {inhibited :: InhibitTab
inhibited = InhibitTab
inhbt'} BotUpEnv
bu_env {activeCoals :: CoalsTab
activeCoals = CoalsTab
actv'}
else forall (f :: * -> *) a. Applicative f => a -> f a
pure CoalsTab
actv'
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ BotUpEnv
bu_env {activeCoals :: CoalsTab
activeCoals = CoalsTab
actv'', inhibit :: InhibitTab
inhibit = InhibitTab
inhbt'}
mkCoalsTabStm InhibitTab
lutab stm :: Stm (Aliases rep)
stm@(Let pat :: Pat (LetDec (Aliases rep))
pat@(Pat [PatElem (LetDec (Aliases rep))
x']) StmAux (ExpDec (Aliases rep))
_ (BasicOp (FlatUpdate VName
x FlatSlice SubExp
_ VName
_elm))) TopdownEnv rep
td_env BotUpEnv
bu_env
| [(VName
_, MemBlock PrimType
_ ShapeBase SubExp
_ VName
m_x IxFun
_)] <- forall aliases.
Pat (aliases, LetDecMem) -> [(VName, ArrayMemBound)]
getArrMemAssoc Pat (LetDec (Aliases rep))
pat = do
let (CoalsTab
actv, InhibitTab
inhbt) = forall rep (inner :: * -> *).
(AliasableRep rep, Op rep ~ MemOp inner rep,
HasMemBlock (Aliases rep)) =>
TopdownEnv rep
-> BotUpEnv -> Stm (Aliases rep) -> (CoalsTab, InhibitTab)
recordMemRefUses TopdownEnv rep
td_env BotUpEnv
bu_env Stm (Aliases rep)
stm
(CoalsTab
actv', InhibitTab
inhbt') =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_x CoalsTab
actv of
Maybe CoalsEntry
Nothing -> (CoalsTab
actv, InhibitTab
inhbt)
Just CoalsEntry
info ->
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall dec. PatElem dec -> VName
patElemName PatElem (LetDec (Aliases rep))
x') (CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
info) of
Maybe Coalesced
Nothing -> (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
actv, InhibitTab
inhbt) VName
m_x
Just (Coalesced CoalescedKind
k mblk :: ArrayMemBound
mblk@(MemBlock PrimType
_ ShapeBase SubExp
_ VName
_ IxFun
x_indfun) FreeVarSubsts
_) ->
case forall a rep.
FreeIn a =>
ScopeTab rep
-> Map VName (PrimExp VName) -> a -> Maybe FreeVarSubsts
freeVarSubstitutions (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) (BotUpEnv -> Map VName (PrimExp VName)
scals BotUpEnv
bu_env) IxFun
x_indfun of
Just FreeVarSubsts
fv_subs
| forall rep. TopdownEnv rep -> VName -> Bool
isInScope TopdownEnv rep
td_env (CoalsEntry -> VName
dstmem CoalsEntry
info) ->
let coal_etry_x :: Coalesced
coal_etry_x = CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced CoalescedKind
k ArrayMemBound
mblk FreeVarSubsts
fv_subs
info' :: CoalsEntry
info' =
CoalsEntry
info
{ vartab :: Map VName Coalesced
vartab =
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
x Coalesced
coal_etry_x forall a b. (a -> b) -> a -> b
$
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (forall dec. PatElem dec -> VName
patElemName PatElem (LetDec (Aliases rep))
x') Coalesced
coal_etry_x (CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
info)
}
in (forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
m_x CoalsEntry
info' CoalsTab
actv, InhibitTab
inhbt)
Maybe FreeVarSubsts
_ ->
(CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
actv, InhibitTab
inhbt) VName
m_x
CoalsTab
actv'' <- forall rep (inner :: * -> *).
Coalesceable rep inner =>
Stm (Aliases rep)
-> InhibitTab
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep CoalsTab
mkCoalsHelper3PatternMatch Stm (Aliases rep)
stm InhibitTab
lutab TopdownEnv rep
td_env {inhibited :: InhibitTab
inhibited = InhibitTab
inhbt'} BotUpEnv
bu_env {activeCoals :: CoalsTab
activeCoals = CoalsTab
actv'}
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ BotUpEnv
bu_env {activeCoals :: CoalsTab
activeCoals = CoalsTab
actv'', inhibit :: InhibitTab
inhibit = InhibitTab
inhbt'}
mkCoalsTabStm InhibitTab
_ (Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
_ (BasicOp Update {})) TopdownEnv rep
_ BotUpEnv
_ =
forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"In ArrayCoalescing.hs, fun mkCoalsTabStm, illegal pattern for in-place update: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Pat (LetDec (Aliases rep))
pat
mkCoalsTabStm InhibitTab
lutab stm :: Stm (Aliases rep)
stm@(Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
aux (Op Op (Aliases rep)
op)) TopdownEnv rep
td_env BotUpEnv
bu_env = do
InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> MemOp inner (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
on_op <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall rep.
ShortCircuitReader rep
-> InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
onOp
BotUpEnv
bu_env' <- InhibitTab
-> Pat (VarAliases, LetDecMem)
-> Certs
-> MemOp inner (Aliases rep)
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep BotUpEnv
on_op InhibitTab
lutab Pat (LetDec (Aliases rep))
pat (forall dec. StmAux dec -> Certs
stmAuxCerts StmAux (ExpDec (Aliases rep))
aux) Op (Aliases rep)
op TopdownEnv rep
td_env BotUpEnv
bu_env
CoalsTab
activeCoals' <- forall rep (inner :: * -> *).
Coalesceable rep inner =>
Stm (Aliases rep)
-> InhibitTab
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep CoalsTab
mkCoalsHelper3PatternMatch Stm (Aliases rep)
stm InhibitTab
lutab TopdownEnv rep
td_env BotUpEnv
bu_env'
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ BotUpEnv
bu_env' {activeCoals :: CoalsTab
activeCoals = CoalsTab
activeCoals'}
mkCoalsTabStm InhibitTab
lutab stm :: Stm (Aliases rep)
stm@(Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
_ Exp (Aliases rep)
e) TopdownEnv rep
td_env BotUpEnv
bu_env = do
let (CoalsTab
activeCoals', InhibitTab
inhibit') = forall rep (inner :: * -> *).
(AliasableRep rep, Op rep ~ MemOp inner rep,
HasMemBlock (Aliases rep)) =>
TopdownEnv rep
-> BotUpEnv -> Stm (Aliases rep) -> (CoalsTab, InhibitTab)
recordMemRefUses TopdownEnv rep
td_env BotUpEnv
bu_env Stm (Aliases rep)
stm
safe_4 :: Bool
safe_4 = forall rep. Exp rep -> Bool
createsNewArrOK Exp (Aliases rep)
e
((CoalsTab
activeCoals'', InhibitTab
inhibit''), CoalsTab
successCoals') =
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (Bool
-> ((CoalsTab, InhibitTab), CoalsTab)
-> (VName, ArrayMemBound)
-> ((CoalsTab, InhibitTab), CoalsTab)
foldfun Bool
safe_4) ((CoalsTab
activeCoals', InhibitTab
inhibit'), BotUpEnv -> CoalsTab
successCoals BotUpEnv
bu_env) (forall aliases.
Pat (aliases, LetDecMem) -> [(VName, ArrayMemBound)]
getArrMemAssoc Pat (LetDec (Aliases rep))
pat)
CoalsTab
activeCoals''' <- forall rep (inner :: * -> *).
Coalesceable rep inner =>
Stm (Aliases rep)
-> InhibitTab
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep CoalsTab
mkCoalsHelper3PatternMatch Stm (Aliases rep)
stm InhibitTab
lutab TopdownEnv rep
td_env BotUpEnv
bu_env {successCoals :: CoalsTab
successCoals = CoalsTab
successCoals', activeCoals :: CoalsTab
activeCoals = CoalsTab
activeCoals''}
forall (f :: * -> *) a. Applicative f => a -> f a
pure BotUpEnv
bu_env {activeCoals :: CoalsTab
activeCoals = CoalsTab
activeCoals''', inhibit :: InhibitTab
inhibit = InhibitTab
inhibit'', successCoals :: CoalsTab
successCoals = CoalsTab
successCoals'}
where
foldfun :: Bool
-> ((CoalsTab, InhibitTab), CoalsTab)
-> (VName, ArrayMemBound)
-> ((CoalsTab, InhibitTab), CoalsTab)
foldfun Bool
safe_4 ((CoalsTab
a_acc, InhibitTab
inhb), CoalsTab
s_acc) (VName
b, MemBlock PrimType
tp ShapeBase SubExp
shp VName
mb IxFun
_b_indfun) =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
mb CoalsTab
a_acc of
Maybe CoalsEntry
Nothing -> ((CoalsTab
a_acc, InhibitTab
inhb), CoalsTab
s_acc)
Just info :: CoalsEntry
info@(CoalsEntry VName
x_mem IxFun
_ Names
_ Map VName Coalesced
vtab Map VName VName
_ MemRefs
_ Certs
certs) ->
let failed :: (CoalsTab, InhibitTab)
failed = (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
a_acc, InhibitTab
inhb) VName
mb
in case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
b Map VName Coalesced
vtab of
Maybe Coalesced
Nothing ->
case forall rep.
HasMemBlock (Aliases rep) =>
TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, IxFun)
getDirAliasedIxfn TopdownEnv rep
td_env CoalsTab
a_acc VName
b of
Maybe (VName, VName, IxFun)
Nothing -> ((CoalsTab, InhibitTab)
failed, CoalsTab
s_acc)
Just (VName
_, VName
_, IxFun
b_indfun') ->
case ( forall a rep.
FreeIn a =>
ScopeTab rep
-> Map VName (PrimExp VName) -> a -> Maybe FreeVarSubsts
freeVarSubstitutions (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) (BotUpEnv -> Map VName (PrimExp VName)
scals BotUpEnv
bu_env) IxFun
b_indfun',
forall a rep.
FreeIn a =>
ScopeTab rep
-> Map VName (PrimExp VName) -> a -> Maybe FreeVarSubsts
freeVarSubstitutions (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) (BotUpEnv -> Map VName (PrimExp VName)
scals BotUpEnv
bu_env) Certs
certs
) of
(Just FreeVarSubsts
fv_subst, Just FreeVarSubsts
fv_subst') ->
let mem_info :: Coalesced
mem_info = CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced CoalescedKind
TransitiveCoal (PrimType -> ShapeBase SubExp -> VName -> IxFun -> ArrayMemBound
MemBlock PrimType
tp ShapeBase SubExp
shp VName
x_mem IxFun
b_indfun') (FreeVarSubsts
fv_subst forall a. Semigroup a => a -> a -> a
<> FreeVarSubsts
fv_subst')
info' :: CoalsEntry
info' = CoalsEntry
info {vartab :: Map VName Coalesced
vartab = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
b Coalesced
mem_info Map VName Coalesced
vtab}
in ((forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
mb CoalsEntry
info' CoalsTab
a_acc, InhibitTab
inhb), CoalsTab
s_acc)
(Maybe FreeVarSubsts, Maybe FreeVarSubsts)
_ -> ((CoalsTab, InhibitTab)
failed, CoalsTab
s_acc)
Just (Coalesced CoalescedKind
k mblk :: ArrayMemBound
mblk@(MemBlock PrimType
_ ShapeBase SubExp
_ VName
_ IxFun
new_indfun) FreeVarSubsts
_) ->
let safe_2 :: Bool
safe_2 = forall rep. TopdownEnv rep -> VName -> Bool
isInScope TopdownEnv rep
td_env VName
x_mem
in case ( forall a rep.
FreeIn a =>
ScopeTab rep
-> Map VName (PrimExp VName) -> a -> Maybe FreeVarSubsts
freeVarSubstitutions (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) (BotUpEnv -> Map VName (PrimExp VName)
scals BotUpEnv
bu_env) IxFun
new_indfun,
forall a rep.
FreeIn a =>
ScopeTab rep
-> Map VName (PrimExp VName) -> a -> Maybe FreeVarSubsts
freeVarSubstitutions (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) (BotUpEnv -> Map VName (PrimExp VName)
scals BotUpEnv
bu_env) Certs
certs
) of
(Just FreeVarSubsts
fv_subst, Just FreeVarSubsts
fv_subst')
| Bool
safe_2 ->
let mem_info :: Coalesced
mem_info = CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced CoalescedKind
k ArrayMemBound
mblk (FreeVarSubsts
fv_subst forall a. Semigroup a => a -> a -> a
<> FreeVarSubsts
fv_subst')
info' :: CoalsEntry
info' = CoalsEntry
info {vartab :: Map VName Coalesced
vartab = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
b Coalesced
mem_info Map VName Coalesced
vtab}
in if Bool
safe_4
then
let (CoalsTab
a_acc', CoalsTab
s_acc') = (CoalsTab, CoalsTab) -> VName -> CoalsEntry -> (CoalsTab, CoalsTab)
markSuccessCoal (CoalsTab
a_acc, CoalsTab
s_acc) VName
mb CoalsEntry
info'
in ((CoalsTab
a_acc', InhibitTab
inhb), CoalsTab
s_acc')
else
((forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
mb CoalsEntry
info' CoalsTab
a_acc, InhibitTab
inhb), CoalsTab
s_acc)
(Maybe FreeVarSubsts, Maybe FreeVarSubsts)
_ -> ((CoalsTab, InhibitTab)
failed, CoalsTab
s_acc)
ixfunToAccessSummary :: IxFun.IxFun (TPrimExp Int64 VName) -> AccessSummary
ixfunToAccessSummary :: IxFun -> AccessSummary
ixfunToAccessSummary (IxFun.IxFun LmadRef
lmad [TPrimExp Int64 VName]
_) = Set LmadRef -> AccessSummary
Set forall a b. (a -> b) -> a -> b
$ forall a. a -> Set a
S.singleton LmadRef
lmad
filterSafetyCond2and5 ::
(HasMemBlock (Aliases rep)) =>
CoalsTab ->
InhibitTab ->
ScalarTab ->
TopdownEnv rep ->
[PatElem (VarAliases, LetDecMem)] ->
(CoalsTab, InhibitTab)
filterSafetyCond2and5 :: forall rep.
HasMemBlock (Aliases rep) =>
CoalsTab
-> InhibitTab
-> Map VName (PrimExp VName)
-> TopdownEnv rep
-> [PatElem (VarAliases, LetDecMem)]
-> (CoalsTab, InhibitTab)
filterSafetyCond2and5 CoalsTab
act_coal InhibitTab
inhb_coal Map VName (PrimExp VName)
scals_env TopdownEnv rep
td_env [PatElem (VarAliases, LetDecMem)]
pes =
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (CoalsTab, InhibitTab)
-> PatElem (VarAliases, LetDecMem) -> (CoalsTab, InhibitTab)
helper (CoalsTab
act_coal, InhibitTab
inhb_coal) [PatElem (VarAliases, LetDecMem)]
pes
where
helper :: (CoalsTab, InhibitTab)
-> PatElem (VarAliases, LetDecMem) -> (CoalsTab, InhibitTab)
helper (CoalsTab
acc, InhibitTab
inhb) PatElem (VarAliases, LetDecMem)
patel = do
case (forall dec. PatElem dec -> VName
patElemName PatElem (VarAliases, LetDecMem)
patel, forall dec. PatElem dec -> dec
patElemDec PatElem (VarAliases, LetDecMem)
patel) of
(VName
b, (VarAliases
_, MemArray PrimType
tp0 ShapeBase SubExp
shp0 NoUniqueness
_ (ArrayIn VName
m_b IxFun
_idxfn_b))) ->
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_b CoalsTab
acc of
Maybe CoalsEntry
Nothing -> (CoalsTab
acc, InhibitTab
inhb)
Just info :: CoalsEntry
info@(CoalsEntry VName
x_mem IxFun
_ Names
_ Map VName Coalesced
vtab Map VName VName
_ MemRefs
_ Certs
certs) ->
let failed :: (CoalsTab, InhibitTab)
failed = (CoalsTab, InhibitTab) -> VName -> (CoalsTab, InhibitTab)
markFailedCoal (CoalsTab
acc, InhibitTab
inhb) VName
m_b
in
if forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((VName -> Names -> Bool
`nameIn` forall a. AliasesOf a => a -> Names
aliasesOf PatElem (VarAliases, LetDecMem)
patel) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. PatElem dec -> VName
patElemName) [PatElem (VarAliases, LetDecMem)]
pes
then (CoalsTab, InhibitTab)
failed
else case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
b Map VName Coalesced
vtab of
Maybe Coalesced
Nothing ->
case forall rep.
HasMemBlock (Aliases rep) =>
TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, IxFun)
getDirAliasedIxfn TopdownEnv rep
td_env CoalsTab
acc VName
b of
Maybe (VName, VName, IxFun)
Nothing -> (CoalsTab, InhibitTab)
failed
Just (VName
_, VName
_, IxFun
b_indfun') ->
case ( forall a rep.
FreeIn a =>
ScopeTab rep
-> Map VName (PrimExp VName) -> a -> Maybe FreeVarSubsts
freeVarSubstitutions (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) Map VName (PrimExp VName)
scals_env IxFun
b_indfun',
forall a rep.
FreeIn a =>
ScopeTab rep
-> Map VName (PrimExp VName) -> a -> Maybe FreeVarSubsts
freeVarSubstitutions (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) Map VName (PrimExp VName)
scals_env Certs
certs
) of
(Just FreeVarSubsts
fv_subst, Just FreeVarSubsts
fv_subst') ->
let mem_info :: Coalesced
mem_info = CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced CoalescedKind
TransitiveCoal (PrimType -> ShapeBase SubExp -> VName -> IxFun -> ArrayMemBound
MemBlock PrimType
tp0 ShapeBase SubExp
shp0 VName
x_mem IxFun
b_indfun') (FreeVarSubsts
fv_subst forall a. Semigroup a => a -> a -> a
<> FreeVarSubsts
fv_subst')
info' :: CoalsEntry
info' = CoalsEntry
info {vartab :: Map VName Coalesced
vartab = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
b Coalesced
mem_info Map VName Coalesced
vtab}
in (forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
m_b CoalsEntry
info' CoalsTab
acc, InhibitTab
inhb)
(Maybe FreeVarSubsts, Maybe FreeVarSubsts)
_ -> (CoalsTab, InhibitTab)
failed
Just (Coalesced CoalescedKind
k (MemBlock PrimType
pt ShapeBase SubExp
shp VName
_ IxFun
new_indfun) FreeVarSubsts
_) ->
let safe_2 :: Bool
safe_2 = forall rep. TopdownEnv rep -> VName -> Bool
isInScope TopdownEnv rep
td_env VName
x_mem
in case ( forall a rep.
FreeIn a =>
ScopeTab rep
-> Map VName (PrimExp VName) -> a -> Maybe FreeVarSubsts
freeVarSubstitutions (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) Map VName (PrimExp VName)
scals_env IxFun
new_indfun,
forall a rep.
FreeIn a =>
ScopeTab rep
-> Map VName (PrimExp VName) -> a -> Maybe FreeVarSubsts
freeVarSubstitutions (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) Map VName (PrimExp VName)
scals_env Certs
certs
) of
(Just FreeVarSubsts
fv_subst, Just FreeVarSubsts
fv_subst')
| Bool
safe_2 ->
let mem_info :: Coalesced
mem_info = CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced CoalescedKind
k (PrimType -> ShapeBase SubExp -> VName -> IxFun -> ArrayMemBound
MemBlock PrimType
pt ShapeBase SubExp
shp VName
x_mem IxFun
new_indfun) (FreeVarSubsts
fv_subst forall a. Semigroup a => a -> a -> a
<> FreeVarSubsts
fv_subst')
info' :: CoalsEntry
info' = CoalsEntry
info {vartab :: Map VName Coalesced
vartab = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
b Coalesced
mem_info Map VName Coalesced
vtab}
in (forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
m_b CoalsEntry
info' CoalsTab
acc, InhibitTab
inhb)
(Maybe FreeVarSubsts, Maybe FreeVarSubsts)
_ -> (CoalsTab, InhibitTab)
failed
(VName, (VarAliases, LetDecMem))
_ -> (CoalsTab
acc, InhibitTab
inhb)
mkCoalsHelper3PatternMatch ::
(Coalesceable rep inner) =>
Stm (Aliases rep) ->
LUTabFun ->
TopdownEnv rep ->
BotUpEnv ->
ShortCircuitM rep CoalsTab
mkCoalsHelper3PatternMatch :: forall rep (inner :: * -> *).
Coalesceable rep inner =>
Stm (Aliases rep)
-> InhibitTab
-> TopdownEnv rep
-> BotUpEnv
-> ShortCircuitM rep CoalsTab
mkCoalsHelper3PatternMatch Stm (Aliases rep)
stm InhibitTab
lutab TopdownEnv rep
td_env BotUpEnv
bu_env = do
Maybe [SSPointInfo]
clst <- forall rep (inner :: * -> *).
Coalesceable rep inner =>
InhibitTab
-> TopdownEnv rep
-> ScopeTab rep
-> Stm (Aliases rep)
-> ShortCircuitM rep (Maybe [SSPointInfo])
genCoalStmtInfo InhibitTab
lutab TopdownEnv rep
td_env (forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) Stm (Aliases rep)
stm
case Maybe [SSPointInfo]
clst of
Maybe [SSPointInfo]
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure CoalsTab
activeCoals_tab
Just [SSPointInfo]
clst' -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl CoalsTab -> SSPointInfo -> CoalsTab
processNewCoalesce CoalsTab
activeCoals_tab [SSPointInfo]
clst'
where
successCoals_tab :: CoalsTab
successCoals_tab = BotUpEnv -> CoalsTab
successCoals BotUpEnv
bu_env
activeCoals_tab :: CoalsTab
activeCoals_tab = BotUpEnv -> CoalsTab
activeCoals BotUpEnv
bu_env
processNewCoalesce :: CoalsTab -> SSPointInfo -> CoalsTab
processNewCoalesce CoalsTab
acc (CoalescedKind
knd, IxFun -> IxFun
alias_fn, VName
x, VName
m_x, IxFun
ind_x, VName
b, VName
m_b, IxFun
_, PrimType
tp_b, ShapeBase SubExp
shp_b, Certs
certs) =
let proper_coals_tab :: CoalsTab
proper_coals_tab = case CoalescedKind
knd of
CoalescedKind
InPlaceCoal -> CoalsTab
activeCoals_tab
CoalescedKind
_ -> CoalsTab
successCoals_tab
(VName
m_yx, IxFun
ind_yx, Names
mem_yx_al, Map VName VName
x_deps, Certs
certs') =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_x CoalsTab
proper_coals_tab of
Maybe CoalsEntry
Nothing ->
(VName
m_x, IxFun -> IxFun
alias_fn IxFun
ind_x, VName -> Names
oneName VName
m_x, forall k a. Map k a
M.empty, forall a. Monoid a => a
mempty)
Just (CoalsEntry VName
m_y IxFun
ind_y Names
y_al Map VName Coalesced
vtab Map VName VName
x_deps0 MemRefs
_ Certs
certs'') ->
let ind :: IxFun
ind = case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x Map VName Coalesced
vtab of
Just (Coalesced CoalescedKind
_ (MemBlock PrimType
_ ShapeBase SubExp
_ VName
_ IxFun
ixf) FreeVarSubsts
_) ->
IxFun
ixf
Maybe Coalesced
Nothing ->
IxFun
ind_y
in (VName
m_y, IxFun -> IxFun
alias_fn IxFun
ind, VName -> Names
oneName VName
m_x forall a. Semigroup a => a -> a -> a
<> Names
y_al, Map VName VName
x_deps0, Certs
certs forall a. Semigroup a => a -> a -> a
<> Certs
certs'')
m_b_aliased_m_yx :: Bool
m_b_aliased_m_yx = forall rep. TopdownEnv rep -> VName -> [VName] -> Bool
areAnyAliased TopdownEnv rep
td_env VName
m_b [VName
m_yx]
in if Bool -> Bool
not Bool
m_b_aliased_m_yx Bool -> Bool -> Bool
&& forall rep. TopdownEnv rep -> VName -> Bool
isInScope TopdownEnv rep
td_env VName
m_yx
then
let mem_info :: Coalesced
mem_info = CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced CoalescedKind
knd (PrimType -> ShapeBase SubExp -> VName -> IxFun -> ArrayMemBound
MemBlock PrimType
tp_b ShapeBase SubExp
shp_b VName
m_yx IxFun
ind_yx) forall k a. Map k a
M.empty
opts' :: Map VName VName
opts' =
if VName
m_yx forall a. Eq a => a -> a -> Bool
== VName
m_x
then forall k a. Map k a
M.empty
else forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
x VName
m_x Map VName VName
x_deps
vtab :: Map VName Coalesced
vtab = forall k a. k -> a -> Map k a
M.singleton VName
b Coalesced
mem_info
mvtab :: Maybe (Map VName Coalesced)
mvtab = forall rep.
HasMemBlock (Aliases rep) =>
TopdownEnv rep
-> Map VName Coalesced -> VName -> Maybe (Map VName Coalesced)
addInvAliasesVarTab TopdownEnv rep
td_env Map VName Coalesced
vtab VName
b
is_inhibited :: Bool
is_inhibited = case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_b forall a b. (a -> b) -> a -> b
$ forall rep. TopdownEnv rep -> InhibitTab
inhibited TopdownEnv rep
td_env of
Just Names
nms -> VName
m_yx VName -> Names -> Bool
`nameIn` Names
nms
Maybe Names
Nothing -> Bool
False
in case (Bool
is_inhibited, Maybe (Map VName Coalesced)
mvtab) of
(Bool
True, Maybe (Map VName Coalesced)
_) -> CoalsTab
acc
(Bool
_, Maybe (Map VName Coalesced)
Nothing) -> CoalsTab
acc
(Bool
_, Just Map VName Coalesced
vtab') ->
let coal_etry :: CoalsEntry
coal_etry =
VName
-> IxFun
-> Names
-> Map VName Coalesced
-> Map VName VName
-> MemRefs
-> Certs
-> CoalsEntry
CoalsEntry
VName
m_yx
IxFun
ind_yx
Names
mem_yx_al
Map VName Coalesced
vtab'
Map VName VName
opts'
forall a. Monoid a => a
mempty
(Certs
certs forall a. Semigroup a => a -> a -> a
<> Certs
certs')
in forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
m_b CoalsEntry
coal_etry CoalsTab
acc
else CoalsTab
acc
type SSPointInfo =
( CoalescedKind,
IxFun -> IxFun,
VName,
VName,
IxFun,
VName,
VName,
IxFun,
PrimType,
Shape,
Certs
)
type GenSSPoint rep op =
LUTabFun ->
TopdownEnv rep ->
ScopeTab rep ->
Pat (VarAliases, LetDecMem) ->
Certs ->
op ->
Maybe [SSPointInfo]
genSSPointInfoSeqMem ::
GenSSPoint SeqMem (Op (Aliases SeqMem))
genSSPointInfoSeqMem :: GenSSPoint SeqMem (Op (Aliases SeqMem))
genSSPointInfoSeqMem InhibitTab
_ TopdownEnv SeqMem
_ ScopeTab SeqMem
_ Pat (VarAliases, LetDecMem)
_ Certs
_ Op (Aliases SeqMem)
_ =
forall a. Maybe a
Nothing
genSSPointInfoSegOp ::
(Coalesceable rep inner) => GenSSPoint rep (SegOp lvl (Aliases rep))
genSSPointInfoSegOp :: forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
GenSSPoint rep (SegOp lvl (Aliases rep))
genSSPointInfoSegOp
InhibitTab
lutab
TopdownEnv rep
td_env
ScopeTab rep
scopetab
(Pat [PatElem VName
dst (VarAliases
_, MemArray PrimType
dst_pt ShapeBase SubExp
_ NoUniqueness
_ (ArrayIn VName
dst_mem IxFun
dst_ixf))])
Certs
certs
(SegMap lvl
_ SegSpace
space [Type]
_ kernel_body :: KernelBody (Aliases rep)
kernel_body@KernelBody {kernelBodyResult :: forall rep. KernelBody rep -> [KernelResult]
kernelBodyResult = [Returns {}]})
| (VName
src, MemBlock PrimType
src_pt ShapeBase SubExp
shp VName
src_mem IxFun
src_ixf) : [(VName, ArrayMemBound)]
_ <-
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Stm (Aliases rep) -> Maybe (VName, ArrayMemBound)
getPotentialMapShortCircuit forall a b. (a -> b) -> a -> b
$
forall rep. Stms rep -> [Stm rep]
stmsToList forall a b. (a -> b) -> a -> b
$
forall rep. KernelBody rep -> Stms rep
kernelBodyStms KernelBody (Aliases rep)
kernel_body =
forall a. a -> Maybe a
Just [(CoalescedKind
MapCoal, forall a. a -> a
id, VName
dst, VName
dst_mem, IxFun
dst_ixf, VName
src, VName
src_mem, IxFun
src_ixf, PrimType
src_pt, ShapeBase SubExp
shp, Certs
certs)]
where
iterators :: [VName]
iterators = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space
frees :: Names
frees = forall a. FreeIn a => a -> Names
freeIn KernelBody (Aliases rep)
kernel_body
getPotentialMapShortCircuit :: Stm (Aliases rep) -> Maybe (VName, ArrayMemBound)
getPotentialMapShortCircuit (Let (Pat [PatElem VName
x LetDec (Aliases rep)
_]) StmAux (ExpDec (Aliases rep))
_ (BasicOp (Index VName
src Slice SubExp
slc)))
| Just [SubExp]
inds <- forall d. Slice d -> Maybe [d]
sliceIndices Slice SubExp
slc,
forall a. Ord a => [a] -> [a]
L.sort [SubExp]
inds forall a. Eq a => a -> a -> Bool
== forall a. Ord a => [a] -> [a]
L.sort (forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
iterators),
Just Names
last_uses <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x InhibitTab
lutab,
VName
src VName -> Names -> Bool
`nameIn` Names
last_uses,
Just memblock :: ArrayMemBound
memblock@(MemBlock PrimType
src_pt ShapeBase SubExp
_ VName
src_mem IxFun
src_ixf) <-
forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
src ScopeTab rep
scopetab,
VName
src_mem VName -> Names -> Bool
`nameIn` Names
last_uses,
VName
src_mem forall k a. Ord k => k -> Map k a -> Bool
`M.member` forall rep. TopdownEnv rep -> AllocTab
alloc TopdownEnv rep
td_env,
VName
src VName -> Names -> Bool
`nameIn` Names
frees,
IxFun
src_ixf forall a. Eq a => a -> a -> Bool
== IxFun
dst_ixf,
PrimType -> Int
primBitSize PrimType
src_pt forall a. Eq a => a -> a -> Bool
== PrimType -> Int
primBitSize PrimType
dst_pt =
forall a. a -> Maybe a
Just (VName
src, ArrayMemBound
memblock)
getPotentialMapShortCircuit Stm (Aliases rep)
_ = forall a. Maybe a
Nothing
genSSPointInfoSegOp InhibitTab
_ TopdownEnv rep
_ ScopeTab rep
_ Pat (VarAliases, LetDecMem)
_ Certs
_ SegOp lvl (Aliases rep)
_ =
forall a. Maybe a
Nothing
genSSPointInfoMemOp ::
GenSSPoint rep (inner (Aliases rep)) ->
GenSSPoint rep (MemOp inner (Aliases rep))
genSSPointInfoMemOp :: forall rep (inner :: * -> *).
GenSSPoint rep (inner (Aliases rep))
-> GenSSPoint rep (MemOp inner (Aliases rep))
genSSPointInfoMemOp GenSSPoint rep (inner (Aliases rep))
onOp InhibitTab
lutab TopdownEnv rep
td_end ScopeTab rep
scopetab Pat (VarAliases, LetDecMem)
pat Certs
certs (Inner inner (Aliases rep)
op) =
GenSSPoint rep (inner (Aliases rep))
onOp InhibitTab
lutab TopdownEnv rep
td_end ScopeTab rep
scopetab Pat (VarAliases, LetDecMem)
pat Certs
certs inner (Aliases rep)
op
genSSPointInfoMemOp GenSSPoint rep (inner (Aliases rep))
_ InhibitTab
_ TopdownEnv rep
_ ScopeTab rep
_ Pat (VarAliases, LetDecMem)
_ Certs
_ MemOp inner (Aliases rep)
_ = forall a. Maybe a
Nothing
genSSPointInfoGPUMem ::
GenSSPoint GPUMem (Op (Aliases GPUMem))
genSSPointInfoGPUMem :: GenSSPoint GPUMem (Op (Aliases GPUMem))
genSSPointInfoGPUMem = forall rep (inner :: * -> *).
GenSSPoint rep (inner (Aliases rep))
-> GenSSPoint rep (MemOp inner (Aliases rep))
genSSPointInfoMemOp forall {rep} {inner :: * -> *} {op :: * -> *}.
(LetDec rep ~ LetDecMem, RetType rep ~ RetTypeMem,
LParamInfo rep ~ LetDecMem, OpC rep ~ MemOp inner,
FParamInfo rep ~ FParamMem, BranchType rep ~ BranchTypeMem,
HasLetDecMem (LetDec rep), ASTRep rep, OpReturns (inner rep),
RephraseOp inner, CanBeAliased inner, CanBeAliased (OpC rep),
AliasedOp (OpC rep (Aliases rep)), HasMemBlock (Aliases rep),
TopDownHelper (inner (Aliases rep))) =>
InhibitTab
-> TopdownEnv rep
-> ScopeTab rep
-> Pat (VarAliases, LetDecMem)
-> Certs
-> HostOp op (Aliases rep)
-> Maybe [SSPointInfo]
f
where
f :: InhibitTab
-> TopdownEnv rep
-> ScopeTab rep
-> Pat (VarAliases, LetDecMem)
-> Certs
-> HostOp op (Aliases rep)
-> Maybe [SSPointInfo]
f InhibitTab
lutab TopdownEnv rep
td_env ScopeTab rep
scopetab Pat (VarAliases, LetDecMem)
pat Certs
certs (GPU.SegOp SegOp SegLevel (Aliases rep)
op) =
forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
GenSSPoint rep (SegOp lvl (Aliases rep))
genSSPointInfoSegOp InhibitTab
lutab TopdownEnv rep
td_env ScopeTab rep
scopetab Pat (VarAliases, LetDecMem)
pat Certs
certs SegOp SegLevel (Aliases rep)
op
f InhibitTab
_ TopdownEnv rep
_ ScopeTab rep
_ Pat (VarAliases, LetDecMem)
_ Certs
_ HostOp op (Aliases rep)
_ = forall a. Maybe a
Nothing
genSSPointInfoMCMem ::
GenSSPoint MCMem (Op (Aliases MCMem))
genSSPointInfoMCMem :: GenSSPoint MCMem (Op (Aliases MCMem))
genSSPointInfoMCMem = forall rep (inner :: * -> *).
GenSSPoint rep (inner (Aliases rep))
-> GenSSPoint rep (MemOp inner (Aliases rep))
genSSPointInfoMemOp forall {rep} {inner :: * -> *} {op :: * -> *}.
(LetDec rep ~ LetDecMem, RetType rep ~ RetTypeMem,
LParamInfo rep ~ LetDecMem, OpC rep ~ MemOp inner,
FParamInfo rep ~ FParamMem, BranchType rep ~ BranchTypeMem,
HasLetDecMem (LetDec rep), ASTRep rep, OpReturns (inner rep),
RephraseOp inner, CanBeAliased inner, CanBeAliased (OpC rep),
AliasedOp (OpC rep (Aliases rep)), HasMemBlock (Aliases rep),
TopDownHelper (inner (Aliases rep))) =>
InhibitTab
-> TopdownEnv rep
-> ScopeTab rep
-> Pat (VarAliases, LetDecMem)
-> Certs
-> MCOp op (Aliases rep)
-> Maybe [SSPointInfo]
f
where
f :: InhibitTab
-> TopdownEnv rep
-> ScopeTab rep
-> Pat (VarAliases, LetDecMem)
-> Certs
-> MCOp op (Aliases rep)
-> Maybe [SSPointInfo]
f InhibitTab
lutab TopdownEnv rep
td_env ScopeTab rep
scopetab Pat (VarAliases, LetDecMem)
pat Certs
certs (MC.ParOp Maybe (SegOp () (Aliases rep))
Nothing SegOp () (Aliases rep)
op) =
forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
GenSSPoint rep (SegOp lvl (Aliases rep))
genSSPointInfoSegOp InhibitTab
lutab TopdownEnv rep
td_env ScopeTab rep
scopetab Pat (VarAliases, LetDecMem)
pat Certs
certs SegOp () (Aliases rep)
op
f InhibitTab
_ TopdownEnv rep
_ ScopeTab rep
_ Pat (VarAliases, LetDecMem)
_ Certs
_ MCOp op (Aliases rep)
_ = forall a. Maybe a
Nothing
genCoalStmtInfo ::
(Coalesceable rep inner) =>
LUTabFun ->
TopdownEnv rep ->
ScopeTab rep ->
Stm (Aliases rep) ->
ShortCircuitM rep (Maybe [SSPointInfo])
genCoalStmtInfo :: forall rep (inner :: * -> *).
Coalesceable rep inner =>
InhibitTab
-> TopdownEnv rep
-> ScopeTab rep
-> Stm (Aliases rep)
-> ShortCircuitM rep (Maybe [SSPointInfo])
genCoalStmtInfo InhibitTab
lutab TopdownEnv rep
_ ScopeTab rep
scopetab (Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
aux (BasicOp (Replicate (Shape []) (Var VName
b))))
| Pat [PatElem VName
x (VarAliases
_, MemArray PrimType
_ ShapeBase SubExp
_ NoUniqueness
_ (ArrayIn VName
m_x IxFun
ind_x))] <- Pat (LetDec (Aliases rep))
pat =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ case (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x InhibitTab
lutab, forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
b ScopeTab rep
scopetab) of
(Just Names
last_uses, Just (MemBlock PrimType
tpb ShapeBase SubExp
shpb VName
m_b IxFun
ind_b)) ->
if VName
b VName -> Names -> Bool
`notNameIn` Names
last_uses
then forall a. Maybe a
Nothing
else forall a. a -> Maybe a
Just [(CoalescedKind
CopyCoal, forall a. a -> a
id, VName
x, VName
m_x, IxFun
ind_x, VName
b, VName
m_b, IxFun
ind_b, PrimType
tpb, ShapeBase SubExp
shpb, forall dec. StmAux dec -> Certs
stmAuxCerts StmAux (ExpDec (Aliases rep))
aux)]
(Maybe Names, Maybe ArrayMemBound)
_ -> forall a. Maybe a
Nothing
genCoalStmtInfo InhibitTab
lutab TopdownEnv rep
_ ScopeTab rep
scopetab (Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
aux (BasicOp (Update Safety
_ VName
x Slice SubExp
slice_x (Var VName
b))))
| Pat [PatElem VName
x' (VarAliases
_, MemArray PrimType
_ ShapeBase SubExp
_ NoUniqueness
_ (ArrayIn VName
m_x IxFun
ind_x))] <- Pat (LetDec (Aliases rep))
pat =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ case (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x' InhibitTab
lutab, forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
b ScopeTab rep
scopetab) of
(Just Names
last_uses, Just (MemBlock PrimType
tpb ShapeBase SubExp
shpb VName
m_b IxFun
ind_b)) ->
if VName
b VName -> Names -> Bool
`notNameIn` Names
last_uses
then forall a. Maybe a
Nothing
else forall a. a -> Maybe a
Just [(CoalescedKind
InPlaceCoal, (IxFun -> Slice SubExp -> IxFun
`updateIndFunSlice` Slice SubExp
slice_x), VName
x, VName
m_x, IxFun
ind_x, VName
b, VName
m_b, IxFun
ind_b, PrimType
tpb, ShapeBase SubExp
shpb, forall dec. StmAux dec -> Certs
stmAuxCerts StmAux (ExpDec (Aliases rep))
aux)]
(Maybe Names, Maybe ArrayMemBound)
_ -> forall a. Maybe a
Nothing
where
updateIndFunSlice :: IxFun -> Slice SubExp -> IxFun
updateIndFunSlice :: IxFun -> Slice SubExp -> IxFun
updateIndFunSlice IxFun
ind_fun Slice SubExp
slc_x =
let slc_x' :: [DimIndex (TPrimExp Int64 VName)]
slc_x' = forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64) forall a b. (a -> b) -> a -> b
$ forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slc_x
in forall num.
(Eq num, IntegralExp num) =>
IxFun num -> Slice num -> IxFun num
IxFun.slice IxFun
ind_fun forall a b. (a -> b) -> a -> b
$ forall d. [DimIndex d] -> Slice d
Slice [DimIndex (TPrimExp Int64 VName)]
slc_x'
genCoalStmtInfo InhibitTab
lutab TopdownEnv rep
_ ScopeTab rep
scopetab (Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
aux (BasicOp (FlatUpdate VName
x FlatSlice SubExp
slice_x VName
b)))
| Pat [PatElem VName
x' (VarAliases
_, MemArray PrimType
_ ShapeBase SubExp
_ NoUniqueness
_ (ArrayIn VName
m_x IxFun
ind_x))] <- Pat (LetDec (Aliases rep))
pat =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ case (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x' InhibitTab
lutab, forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
b ScopeTab rep
scopetab) of
(Just Names
last_uses, Just (MemBlock PrimType
tpb ShapeBase SubExp
shpb VName
m_b IxFun
ind_b)) ->
if VName
b VName -> Names -> Bool
`notNameIn` Names
last_uses
then forall a. Maybe a
Nothing
else forall a. a -> Maybe a
Just [(CoalescedKind
InPlaceCoal, (IxFun -> FlatSlice SubExp -> IxFun
`updateIndFunSlice` FlatSlice SubExp
slice_x), VName
x, VName
m_x, IxFun
ind_x, VName
b, VName
m_b, IxFun
ind_b, PrimType
tpb, ShapeBase SubExp
shpb, forall dec. StmAux dec -> Certs
stmAuxCerts StmAux (ExpDec (Aliases rep))
aux)]
(Maybe Names, Maybe ArrayMemBound)
_ -> forall a. Maybe a
Nothing
where
updateIndFunSlice :: IxFun -> FlatSlice SubExp -> IxFun
updateIndFunSlice :: IxFun -> FlatSlice SubExp -> IxFun
updateIndFunSlice IxFun
ind_fun (FlatSlice SubExp
offset [FlatDimIndex SubExp]
dims) =
forall num.
(Eq num, IntegralExp num) =>
IxFun num -> FlatSlice num -> IxFun num
IxFun.flatSlice IxFun
ind_fun forall a b. (a -> b) -> a -> b
$ forall d. d -> [FlatDimIndex d] -> FlatSlice d
FlatSlice (SubExp -> TPrimExp Int64 VName
pe64 SubExp
offset) forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64) [FlatDimIndex SubExp]
dims
genCoalStmtInfo InhibitTab
lutab TopdownEnv rep
_ ScopeTab rep
scopetab (Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
aux (BasicOp (Concat Int
concat_dim (VName
b0 :| [VName]
bs) SubExp
_)))
| Pat [PatElem VName
x (VarAliases
_, MemArray PrimType
_ ShapeBase SubExp
_ NoUniqueness
_ (ArrayIn VName
m_x IxFun
ind_x))] <- Pat (LetDec (Aliases rep))
pat =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x InhibitTab
lutab of
Maybe Names
Nothing -> forall a. Maybe a
Nothing
Just Names
last_uses ->
let zero :: TPrimExp Int64 VName
zero = SubExp -> TPrimExp Int64 VName
pe64 forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
markConcatParts :: ([SSPointInfo], TPrimExp Int64 VName, Bool)
-> VName -> ([SSPointInfo], TPrimExp Int64 VName, Bool)
markConcatParts ([SSPointInfo]
acc, TPrimExp Int64 VName
offs, Bool
succ0) VName
b =
if Bool -> Bool
not Bool
succ0
then ([SSPointInfo]
acc, TPrimExp Int64 VName
offs, Bool
succ0)
else case forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
b ScopeTab rep
scopetab of
Just (MemBlock PrimType
tpb shpb :: ShapeBase SubExp
shpb@(Shape dims :: [SubExp]
dims@(SubExp
_ : [SubExp]
_)) VName
m_b IxFun
ind_b)
| Just SubExp
d <- forall int a. Integral int => int -> [a] -> Maybe a
maybeNth Int
concat_dim [SubExp]
dims ->
let offs' :: TPrimExp Int64 VName
offs' = TPrimExp Int64 VName
offs forall a. Num a => a -> a -> a
+ SubExp -> TPrimExp Int64 VName
pe64 SubExp
d
in if VName
b VName -> Names -> Bool
`nameIn` Names
last_uses
then
let slc :: Slice (TPrimExp Int64 VName)
slc =
forall d. [DimIndex d] -> Slice d
Slice forall a b. (a -> b) -> a -> b
$
forall a b. (a -> b) -> [a] -> [b]
map (forall d. Num d => d -> d -> DimIndex d
unitSlice TPrimExp Int64 VName
zero forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> TPrimExp Int64 VName
pe64) (forall a. Int -> [a] -> [a]
take Int
concat_dim [SubExp]
dims)
forall a. Semigroup a => a -> a -> a
<> [forall d. Num d => d -> d -> DimIndex d
unitSlice TPrimExp Int64 VName
offs (SubExp -> TPrimExp Int64 VName
pe64 SubExp
d)]
forall a. Semigroup a => a -> a -> a
<> forall a b. (a -> b) -> [a] -> [b]
map (forall d. Num d => d -> d -> DimIndex d
unitSlice TPrimExp Int64 VName
zero forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> TPrimExp Int64 VName
pe64) (forall a. Int -> [a] -> [a]
drop (Int
concat_dim forall a. Num a => a -> a -> a
+ Int
1) [SubExp]
dims)
in ( [SSPointInfo]
acc forall a. [a] -> [a] -> [a]
++ [(CoalescedKind
ConcatCoal, (forall num.
(Eq num, IntegralExp num) =>
IxFun num -> Slice num -> IxFun num
`IxFun.slice` Slice (TPrimExp Int64 VName)
slc), VName
x, VName
m_x, IxFun
ind_x, VName
b, VName
m_b, IxFun
ind_b, PrimType
tpb, ShapeBase SubExp
shpb, forall dec. StmAux dec -> Certs
stmAuxCerts StmAux (ExpDec (Aliases rep))
aux)],
TPrimExp Int64 VName
offs',
Bool
True
)
else ([SSPointInfo]
acc, TPrimExp Int64 VName
offs', Bool
True)
Maybe ArrayMemBound
_ -> ([SSPointInfo]
acc, TPrimExp Int64 VName
offs, Bool
False)
([SSPointInfo]
res, TPrimExp Int64 VName
_, Bool
_) = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ([SSPointInfo], TPrimExp Int64 VName, Bool)
-> VName -> ([SSPointInfo], TPrimExp Int64 VName, Bool)
markConcatParts ([], TPrimExp Int64 VName
zero, Bool
True) (VName
b0 forall a. a -> [a] -> [a]
: [VName]
bs)
in if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [SSPointInfo]
res then forall a. Maybe a
Nothing else forall a. a -> Maybe a
Just [SSPointInfo]
res
genCoalStmtInfo InhibitTab
lutab TopdownEnv rep
td_env ScopeTab rep
scopetab (Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
aux (Op Op (Aliases rep)
op)) = do
InhibitTab
-> TopdownEnv rep
-> ScopeTab rep
-> Pat (VarAliases, LetDecMem)
-> Certs
-> OpC rep (Aliases rep)
-> Maybe [SSPointInfo]
ss_op <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall rep.
ShortCircuitReader rep
-> InhibitTab
-> TopdownEnv rep
-> ScopeTab rep
-> Pat (VarAliases, LetDecMem)
-> Certs
-> Op (Aliases rep)
-> Maybe [SSPointInfo]
ssPointFromOp
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ InhibitTab
-> TopdownEnv rep
-> ScopeTab rep
-> Pat (VarAliases, LetDecMem)
-> Certs
-> OpC rep (Aliases rep)
-> Maybe [SSPointInfo]
ss_op InhibitTab
lutab TopdownEnv rep
td_env ScopeTab rep
scopetab Pat (LetDec (Aliases rep))
pat (forall dec. StmAux dec -> Certs
stmAuxCerts StmAux (ExpDec (Aliases rep))
aux) Op (Aliases rep)
op
genCoalStmtInfo InhibitTab
_ TopdownEnv rep
_ ScopeTab rep
_ Stm (Aliases rep)
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
data MemBodyResult = MemBodyResult
{ MemBodyResult -> VName
patMem :: VName,
MemBodyResult -> VName
_patName :: VName,
MemBodyResult -> VName
bodyName :: VName,
MemBodyResult -> VName
bodyMem :: VName
}
findMemBodyResult ::
(HasMemBlock (Aliases rep)) =>
CoalsTab ->
ScopeTab rep ->
[PatElem (VarAliases, LetDecMem)] ->
Body (Aliases rep) ->
[MemBodyResult]
findMemBodyResult :: forall rep.
HasMemBlock (Aliases rep) =>
CoalsTab
-> ScopeTab rep
-> [PatElem (VarAliases, LetDecMem)]
-> Body (Aliases rep)
-> [MemBodyResult]
findMemBodyResult CoalsTab
activeCoals_tab ScopeTab rep
scope_env [PatElem (VarAliases, LetDecMem)]
patelms Body (Aliases rep)
bdy =
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
(PatElem (VarAliases, LetDecMem), SubExp) -> Maybe MemBodyResult
findMemBodyResult'
(forall a b. [a] -> [b] -> [(a, b)]
zip [PatElem (VarAliases, LetDecMem)]
patelms forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp forall a b. (a -> b) -> a -> b
$ forall rep. Body rep -> Result
bodyResult Body (Aliases rep)
bdy)
where
scope_env' :: ScopeTab rep
scope_env' = ScopeTab rep
scope_env forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf (forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
bdy)
findMemBodyResult' :: (PatElem (VarAliases, LetDecMem), SubExp) -> Maybe MemBodyResult
findMemBodyResult' (PatElem (VarAliases, LetDecMem)
patel, SubExp
se_r) =
case (forall dec. PatElem dec -> VName
patElemName PatElem (VarAliases, LetDecMem)
patel, forall dec. PatElem dec -> dec
patElemDec PatElem (VarAliases, LetDecMem)
patel, SubExp
se_r) of
(VName
b, (VarAliases
_, MemArray PrimType
_ ShapeBase SubExp
_ NoUniqueness
_ (ArrayIn VName
m_b IxFun
_)), Var VName
r) ->
case forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
r ScopeTab rep
scope_env' of
Maybe ArrayMemBound
Nothing -> forall a. Maybe a
Nothing
Just (MemBlock PrimType
_ ShapeBase SubExp
_ VName
m_r IxFun
_) ->
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_b CoalsTab
activeCoals_tab of
Maybe CoalsEntry
Nothing -> forall a. Maybe a
Nothing
Just CoalsEntry
coal_etry ->
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
b (CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
coal_etry) of
Maybe Coalesced
Nothing -> forall a. Maybe a
Nothing
Just Coalesced
_ -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ VName -> VName -> VName -> VName -> MemBodyResult
MemBodyResult VName
m_b VName
b VName
r VName
m_r
(VName, (VarAliases, LetDecMem), SubExp)
_ -> forall a. Maybe a
Nothing
transferCoalsToBody ::
M.Map VName (TPrimExp Int64 VName) ->
CoalsTab ->
MemBodyResult ->
CoalsTab
transferCoalsToBody :: FreeVarSubsts -> CoalsTab -> MemBodyResult -> CoalsTab
transferCoalsToBody FreeVarSubsts
exist_subs CoalsTab
activeCoals_tab (MemBodyResult VName
m_b VName
b VName
r VName
m_r)
|
Just CoalsEntry
etry <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_b CoalsTab
activeCoals_tab,
Just (Coalesced CoalescedKind
knd (MemBlock PrimType
btp ShapeBase SubExp
shp VName
_ IxFun
ind_b) FreeVarSubsts
subst_b) <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
b forall a b. (a -> b) -> a -> b
$ CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
etry =
let ind_r :: IxFun
ind_r = forall {k} a (t :: k).
Ord a =>
Map a (TPrimExp t a)
-> IxFun (TPrimExp t a) -> IxFun (TPrimExp t a)
IxFun.substituteInIxFun FreeVarSubsts
exist_subs IxFun
ind_b
subst_r :: FreeVarSubsts
subst_r = forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union FreeVarSubsts
exist_subs FreeVarSubsts
subst_b
mem_info :: Coalesced
mem_info = CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced CoalescedKind
knd (PrimType -> ShapeBase SubExp -> VName -> IxFun -> ArrayMemBound
MemBlock PrimType
btp ShapeBase SubExp
shp (CoalsEntry -> VName
dstmem CoalsEntry
etry) IxFun
ind_r) FreeVarSubsts
subst_r
in if VName
m_r forall a. Eq a => a -> a -> Bool
== VName
m_b
then
let etry' :: CoalsEntry
etry' =
CoalsEntry
etry
{ optdeps :: Map VName VName
optdeps = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
b VName
m_b (CoalsEntry -> Map VName VName
optdeps CoalsEntry
etry),
vartab :: Map VName Coalesced
vartab = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
r Coalesced
mem_info (CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
etry)
}
in forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
m_r CoalsEntry
etry' CoalsTab
activeCoals_tab
else
let opts_x_new :: Map VName VName
opts_x_new = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
r VName
m_r (CoalsEntry -> Map VName VName
optdeps CoalsEntry
etry)
coal_etry :: CoalsEntry
coal_etry =
CoalsEntry
etry
{ vartab :: Map VName Coalesced
vartab = forall k a. k -> a -> Map k a
M.singleton VName
r Coalesced
mem_info,
optdeps :: Map VName VName
optdeps = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
b VName
m_b (CoalsEntry -> Map VName VName
optdeps CoalsEntry
etry)
}
in forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
m_b (CoalsEntry
etry {optdeps :: Map VName VName
optdeps = Map VName VName
opts_x_new}) forall a b. (a -> b) -> a -> b
$
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
m_r CoalsEntry
coal_etry CoalsTab
activeCoals_tab
| Bool
otherwise = forall a. HasCallStack => String -> a
error String
"Impossible"
mkSubsTab ::
Pat (aliases, LetDecMem) ->
[SubExp] ->
M.Map VName (TPrimExp Int64 VName)
mkSubsTab :: forall aliases.
Pat (aliases, LetDecMem) -> [SubExp] -> FreeVarSubsts
mkSubsTab Pat (aliases, LetDecMem)
pat [SubExp]
res =
let pat_elms :: [PatElem (aliases, LetDecMem)]
pat_elms = forall dec. Pat dec -> [PatElem dec]
patElems Pat (aliases, LetDecMem)
pat
in forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe forall {a} {d} {u} {ret}.
(PatElem (a, MemInfo d u ret), SubExp)
-> Maybe (VName, TPrimExp Int64 VName)
mki64subst forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [PatElem (aliases, LetDecMem)]
pat_elms [SubExp]
res
where
mki64subst :: (PatElem (a, MemInfo d u ret), SubExp)
-> Maybe (VName, TPrimExp Int64 VName)
mki64subst (PatElem (a, MemInfo d u ret)
a, Var VName
v)
| (a
_, MemPrim (IntType IntType
Int64)) <- forall dec. PatElem dec -> dec
patElemDec PatElem (a, MemInfo d u ret)
a = forall a. a -> Maybe a
Just (forall dec. PatElem dec -> VName
patElemName PatElem (a, MemInfo d u ret)
a, forall a. a -> TPrimExp Int64 a
le64 VName
v)
mki64subst (PatElem (a, MemInfo d u ret)
a, se :: SubExp
se@(Constant (IntValue (Int64Value Int64
_)))) = forall a. a -> Maybe a
Just (forall dec. PatElem dec -> VName
patElemName PatElem (a, MemInfo d u ret)
a, SubExp -> TPrimExp Int64 VName
pe64 SubExp
se)
mki64subst (PatElem (a, MemInfo d u ret), SubExp)
_ = forall a. Maybe a
Nothing
computeScalarTable ::
(Coalesceable rep inner) =>
ScopeTab rep ->
Stm (Aliases rep) ->
ScalarTableM rep (M.Map VName (PrimExp VName))
computeScalarTable :: forall rep (inner :: * -> *).
Coalesceable rep inner =>
ScopeTab rep
-> Stm (Aliases rep)
-> ScalarTableM rep (Map VName (PrimExp VName))
computeScalarTable ScopeTab rep
scope_table (Let (Pat [PatElem (LetDec (Aliases rep))
pe]) StmAux (ExpDec (Aliases rep))
_ Exp (Aliases rep)
e)
| Just PrimExp VName
primexp <- forall (m :: * -> *) rep v.
(MonadFail m, RepTypes rep) =>
(VName -> m (PrimExp v)) -> Exp rep -> m (PrimExp v)
primExpFromExp (forall rep.
AliasableRep rep =>
ScopeTab rep
-> Map VName (PrimExp VName) -> VName -> Maybe (PrimExp VName)
vnameToPrimExp ScopeTab rep
scope_table forall a. Monoid a => a
mempty) Exp (Aliases rep)
e =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton (forall dec. PatElem dec -> VName
patElemName PatElem (LetDec (Aliases rep))
pe) PrimExp VName
primexp
computeScalarTable ScopeTab rep
scope_table (Let Pat (LetDec (Aliases rep))
_ StmAux (ExpDec (Aliases rep))
_ (Loop [(FParam (Aliases rep), SubExp)]
loop_inits LoopForm
loop_form Body (Aliases rep)
body)) =
forall (m :: * -> *) b a.
(Monad m, Monoid b) =>
(a -> m b) -> [a] -> m b
concatMapM
( forall rep (inner :: * -> *).
Coalesceable rep inner =>
ScopeTab rep
-> Stm (Aliases rep)
-> ScalarTableM rep (Map VName (PrimExp VName))
computeScalarTable forall a b. (a -> b) -> a -> b
$
ScopeTab rep
scope_table
forall a. Semigroup a => a -> a -> a
<> forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(FParam (Aliases rep), SubExp)]
loop_inits)
forall a. Semigroup a => a -> a -> a
<> forall rep. LoopForm -> Scope rep
scopeOfLoopForm LoopForm
loop_form
forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf (forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
body)
)
(forall rep. Stms rep -> [Stm rep]
stmsToList forall a b. (a -> b) -> a -> b
$ forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
body)
computeScalarTable ScopeTab rep
scope_table (Let Pat (LetDec (Aliases rep))
_ StmAux (ExpDec (Aliases rep))
_ (Match [SubExp]
_ [Case (Body (Aliases rep))]
cases Body (Aliases rep)
body MatchDec (BranchType (Aliases rep))
_)) = do
Map VName (PrimExp VName)
body_tab <- forall (m :: * -> *) b a.
(Monad m, Monoid b) =>
(a -> m b) -> [a] -> m b
concatMapM (forall rep (inner :: * -> *).
Coalesceable rep inner =>
ScopeTab rep
-> Stm (Aliases rep)
-> ScalarTableM rep (Map VName (PrimExp VName))
computeScalarTable forall a b. (a -> b) -> a -> b
$ ScopeTab rep
scope_table forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf (forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
body)) (forall rep. Stms rep -> [Stm rep]
stmsToList forall a b. (a -> b) -> a -> b
$ forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
body)
Map VName (PrimExp VName)
cases_tab <-
forall (m :: * -> *) b a.
(Monad m, Monoid b) =>
(a -> m b) -> [a] -> m b
concatMapM
( \(Case [Maybe PrimValue]
_ Body (Aliases rep)
b) ->
forall (m :: * -> *) b a.
(Monad m, Monoid b) =>
(a -> m b) -> [a] -> m b
concatMapM
(forall rep (inner :: * -> *).
Coalesceable rep inner =>
ScopeTab rep
-> Stm (Aliases rep)
-> ScalarTableM rep (Map VName (PrimExp VName))
computeScalarTable forall a b. (a -> b) -> a -> b
$ ScopeTab rep
scope_table forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf (forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
b))
( forall rep. Stms rep -> [Stm rep]
stmsToList forall a b. (a -> b) -> a -> b
$
forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
body
)
)
[Case (Body (Aliases rep))]
cases
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Map VName (PrimExp VName)
body_tab forall a. Semigroup a => a -> a -> a
<> Map VName (PrimExp VName)
cases_tab
computeScalarTable ScopeTab rep
scope_table (Let Pat (LetDec (Aliases rep))
_ StmAux (ExpDec (Aliases rep))
_ (Op Op (Aliases rep)
op)) = do
ScopeTab rep
-> MemOp inner (Aliases rep)
-> ReaderT
(ComputeScalarTableOnOp rep) Identity (Map VName (PrimExp VName))
on_op <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall rep.
ComputeScalarTableOnOp rep
-> ComputeScalarTable rep (Op (Aliases rep))
scalarTableOnOp
ScopeTab rep
-> MemOp inner (Aliases rep)
-> ReaderT
(ComputeScalarTableOnOp rep) Identity (Map VName (PrimExp VName))
on_op ScopeTab rep
scope_table Op (Aliases rep)
op
computeScalarTable ScopeTab rep
_ Stm (Aliases rep)
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty
computeScalarTableMemOp ::
ComputeScalarTable rep (inner (Aliases rep)) -> ComputeScalarTable rep (MemOp inner (Aliases rep))
computeScalarTableMemOp :: forall rep (inner :: * -> *).
ComputeScalarTable rep (inner (Aliases rep))
-> ComputeScalarTable rep (MemOp inner (Aliases rep))
computeScalarTableMemOp ComputeScalarTable rep (inner (Aliases rep))
_ ScopeTab rep
_ (Alloc SubExp
_ Space
_) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty
computeScalarTableMemOp ComputeScalarTable rep (inner (Aliases rep))
onInner ScopeTab rep
scope_table (Inner inner (Aliases rep)
op) = ComputeScalarTable rep (inner (Aliases rep))
onInner ScopeTab rep
scope_table inner (Aliases rep)
op
computeScalarTableSegOp ::
(Coalesceable rep inner) =>
ComputeScalarTable rep (GPU.SegOp lvl (Aliases rep))
computeScalarTableSegOp :: forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
ComputeScalarTable rep (SegOp lvl (Aliases rep))
computeScalarTableSegOp ScopeTab rep
scope_table SegOp lvl (Aliases rep)
segop = do
forall (m :: * -> *) b a.
(Monad m, Monoid b) =>
(a -> m b) -> [a] -> m b
concatMapM
( forall rep (inner :: * -> *).
Coalesceable rep inner =>
ScopeTab rep
-> Stm (Aliases rep)
-> ScalarTableM rep (Map VName (PrimExp VName))
computeScalarTable forall a b. (a -> b) -> a -> b
$
ScopeTab rep
scope_table
forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf (forall rep. KernelBody rep -> Stms rep
kernelBodyStms forall a b. (a -> b) -> a -> b
$ forall lvl rep. SegOp lvl rep -> KernelBody rep
segBody SegOp lvl (Aliases rep)
segop)
forall a. Semigroup a => a -> a -> a
<> forall rep. SegSpace -> Scope rep
scopeOfSegSpace (forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp lvl (Aliases rep)
segop)
)
(forall rep. Stms rep -> [Stm rep]
stmsToList forall a b. (a -> b) -> a -> b
$ forall rep. KernelBody rep -> Stms rep
kernelBodyStms forall a b. (a -> b) -> a -> b
$ forall lvl rep. SegOp lvl rep -> KernelBody rep
segBody SegOp lvl (Aliases rep)
segop)
computeScalarTableGPUMem ::
ComputeScalarTable GPUMem (GPU.HostOp NoOp (Aliases GPUMem))
computeScalarTableGPUMem :: ComputeScalarTable GPUMem (HostOp NoOp (Aliases GPUMem))
computeScalarTableGPUMem ScopeTab GPUMem
scope_table (GPU.SegOp SegOp SegLevel (Aliases GPUMem)
segop) =
forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
ComputeScalarTable rep (SegOp lvl (Aliases rep))
computeScalarTableSegOp ScopeTab GPUMem
scope_table SegOp SegLevel (Aliases GPUMem)
segop
computeScalarTableGPUMem ScopeTab GPUMem
_ (GPU.SizeOp SizeOp
_) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty
computeScalarTableGPUMem ScopeTab GPUMem
_ (GPU.OtherOp NoOp (Aliases GPUMem)
NoOp) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty
computeScalarTableGPUMem ScopeTab GPUMem
scope_table (GPU.GPUBody [Type]
_ Body (Aliases GPUMem)
body) =
forall (m :: * -> *) b a.
(Monad m, Monoid b) =>
(a -> m b) -> [a] -> m b
concatMapM
(forall rep (inner :: * -> *).
Coalesceable rep inner =>
ScopeTab rep
-> Stm (Aliases rep)
-> ScalarTableM rep (Map VName (PrimExp VName))
computeScalarTable forall a b. (a -> b) -> a -> b
$ ScopeTab GPUMem
scope_table forall a. Semigroup a => a -> a -> a
<> forall rep a. Scoped rep a => a -> Scope rep
scopeOf (forall rep. Body rep -> Stms rep
bodyStms Body (Aliases GPUMem)
body))
(forall rep. Stms rep -> [Stm rep]
stmsToList forall a b. (a -> b) -> a -> b
$ forall rep. Body rep -> Stms rep
bodyStms Body (Aliases GPUMem)
body)
computeScalarTableMCMem ::
ComputeScalarTable MCMem (MC.MCOp NoOp (Aliases MCMem))
computeScalarTableMCMem :: ComputeScalarTable MCMem (MCOp NoOp (Aliases MCMem))
computeScalarTableMCMem ScopeTab MCMem
_ (MC.OtherOp NoOp (Aliases MCMem)
NoOp) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty
computeScalarTableMCMem ScopeTab MCMem
scope_table (MC.ParOp Maybe (SegOp () (Aliases MCMem))
par_op SegOp () (Aliases MCMem)
segop) =
forall a. Semigroup a => a -> a -> a
(<>)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty) (forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
ComputeScalarTable rep (SegOp lvl (Aliases rep))
computeScalarTableSegOp ScopeTab MCMem
scope_table) Maybe (SegOp () (Aliases MCMem))
par_op
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall rep (inner :: * -> *) lvl.
Coalesceable rep inner =>
ComputeScalarTable rep (SegOp lvl (Aliases rep))
computeScalarTableSegOp ScopeTab MCMem
scope_table SegOp () (Aliases MCMem)
segop
filterMapM1 :: (Eq k, Monad m) => (v -> m Bool) -> M.Map k v -> m (M.Map k v)
filterMapM1 :: forall k (m :: * -> *) v.
(Eq k, Monad m) =>
(v -> m Bool) -> Map k v -> m (Map k v)
filterMapM1 v -> m Bool
f Map k v
m = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall k a. Eq k => [(k, a)] -> Map k a
M.fromAscList forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM (v -> m Bool
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toAscList Map k v
m