{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}
module Futhark.Optimise.ArrayShortCircuiting.TopdownAnalysis
( TopdownEnv (..),
ScopeTab,
TopDownHelper,
InhibitTab,
updateTopdownEnv,
updateTopdownEnvLoop,
getDirAliasedIxfn,
getDirAliasedIxfn',
addInvAliasesVarTab,
areAnyAliased,
isInScope,
nonNegativesInPat,
)
where
import Data.Map.Strict qualified as M
import Data.Maybe
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.Optimise.ArrayShortCircuiting.DataStructs
type DirAlias = IxFun -> Maybe IxFun
type InvAlias = Maybe (IxFun -> IxFun)
type VarAliasTab = M.Map VName (VName, DirAlias, InvAlias)
type MemAliasTab = M.Map VName Names
data TopdownEnv rep = TopdownEnv
{
forall rep. TopdownEnv rep -> AllocTab
alloc :: AllocTab,
forall rep. TopdownEnv rep -> ScopeTab rep
scope :: ScopeTab rep,
forall rep. TopdownEnv rep -> InhibitTab
inhibited :: InhibitTab,
forall rep. TopdownEnv rep -> VarAliasTab
v_alias :: VarAliasTab,
forall rep. TopdownEnv rep -> InhibitTab
m_alias :: MemAliasTab,
forall rep. TopdownEnv rep -> Names
nonNegatives :: Names,
forall rep. TopdownEnv rep -> Map VName (PrimExp VName)
scalarTable :: M.Map VName (PrimExp VName),
forall rep. TopdownEnv rep -> [(VName, PrimExp VName)]
knownLessThan :: [(VName, PrimExp VName)],
forall rep. TopdownEnv rep -> [SubExp]
td_asserts :: [SubExp]
}
isInScope :: TopdownEnv rep -> VName -> Bool
isInScope :: forall rep. TopdownEnv rep -> VName -> Bool
isInScope TopdownEnv rep
td_env VName
m = VName
m VName -> Map VName (NameInfo (Aliases rep)) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` TopdownEnv rep -> Map VName (NameInfo (Aliases rep))
forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env
getDirAliasFromExp :: Exp (Aliases rep) -> Maybe (VName, DirAlias)
getDirAliasFromExp :: forall rep. Exp (Aliases rep) -> Maybe (VName, DirAlias)
getDirAliasFromExp (BasicOp (SubExp (Var VName
x))) = (VName, DirAlias) -> Maybe (VName, DirAlias)
forall a. a -> Maybe a
Just (VName
x, DirAlias
forall a. a -> Maybe a
Just)
getDirAliasFromExp (BasicOp (Opaque OpaqueOp
_ (Var VName
x))) = (VName, DirAlias) -> Maybe (VName, DirAlias)
forall a. a -> Maybe a
Just (VName
x, DirAlias
forall a. a -> Maybe a
Just)
getDirAliasFromExp (BasicOp (Reshape ReshapeKind
ReshapeCoerce Shape
shp VName
x)) =
(VName, DirAlias) -> Maybe (VName, DirAlias)
forall a. a -> Maybe a
Just (VName
x, DirAlias
forall a. a -> Maybe a
Just DirAlias -> (IxFun -> IxFun) -> DirAlias
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IxFun -> Shape (TPrimExp Int64 VName) -> IxFun
forall num.
(Eq num, IntegralExp num) =>
IxFun num -> Shape num -> IxFun num
`IxFun.coerce` ShapeBase (TPrimExp Int64 VName) -> Shape (TPrimExp Int64 VName)
forall d. ShapeBase d -> [d]
shapeDims ((SubExp -> TPrimExp Int64 VName)
-> Shape -> ShapeBase (TPrimExp Int64 VName)
forall a b. (a -> b) -> ShapeBase a -> ShapeBase b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 Shape
shp)))
getDirAliasFromExp (BasicOp (Reshape ReshapeKind
ReshapeArbitrary Shape
shp VName
x)) =
(VName, DirAlias) -> Maybe (VName, DirAlias)
forall a. a -> Maybe a
Just (VName
x, (IxFun -> Shape (TPrimExp Int64 VName) -> Maybe IxFun
forall num.
(Eq num, IntegralExp num) =>
IxFun num -> Shape num -> Maybe (IxFun num)
`IxFun.reshape` ShapeBase (TPrimExp Int64 VName) -> Shape (TPrimExp Int64 VName)
forall d. ShapeBase d -> [d]
shapeDims ((SubExp -> TPrimExp Int64 VName)
-> Shape -> ShapeBase (TPrimExp Int64 VName)
forall a b. (a -> b) -> ShapeBase a -> ShapeBase b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 Shape
shp)))
getDirAliasFromExp (BasicOp (Rearrange [Int]
_ VName
_)) =
Maybe (VName, DirAlias)
forall a. Maybe a
Nothing
getDirAliasFromExp (BasicOp (Index VName
x Slice SubExp
slc)) =
(VName, DirAlias) -> Maybe (VName, DirAlias)
forall a. a -> Maybe a
Just (VName
x, DirAlias
forall a. a -> Maybe a
Just DirAlias -> (IxFun -> IxFun) -> DirAlias
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IxFun -> Slice (TPrimExp Int64 VName) -> IxFun
forall num.
(Eq num, IntegralExp num) =>
IxFun num -> Slice num -> IxFun num
`IxFun.slice` ([DimIndex (TPrimExp Int64 VName)] -> Slice (TPrimExp Int64 VName)
forall d. [DimIndex d] -> Slice d
Slice ([DimIndex (TPrimExp Int64 VName)] -> Slice (TPrimExp Int64 VName))
-> [DimIndex (TPrimExp Int64 VName)]
-> Slice (TPrimExp Int64 VName)
forall a b. (a -> b) -> a -> b
$ (DimIndex SubExp -> DimIndex (TPrimExp Int64 VName))
-> [DimIndex SubExp] -> [DimIndex (TPrimExp Int64 VName)]
forall a b. (a -> b) -> [a] -> [b]
map ((SubExp -> TPrimExp Int64 VName)
-> DimIndex SubExp -> DimIndex (TPrimExp Int64 VName)
forall a b. (a -> b) -> DimIndex a -> DimIndex b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64) ([DimIndex SubExp] -> [DimIndex (TPrimExp Int64 VName)])
-> [DimIndex SubExp] -> [DimIndex (TPrimExp Int64 VName)]
forall a b. (a -> b) -> a -> b
$ Slice SubExp -> [DimIndex SubExp]
forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slc)))
getDirAliasFromExp (BasicOp (Update Safety
_ VName
x Slice SubExp
_ SubExp
_elm)) = (VName, DirAlias) -> Maybe (VName, DirAlias)
forall a. a -> Maybe a
Just (VName
x, DirAlias
forall a. a -> Maybe a
Just)
getDirAliasFromExp (BasicOp (FlatIndex VName
x (FlatSlice SubExp
offset [FlatDimIndex SubExp]
idxs))) =
(VName, DirAlias) -> Maybe (VName, DirAlias)
forall a. a -> Maybe a
Just
( VName
x,
DirAlias
forall a. a -> Maybe a
Just DirAlias -> (IxFun -> IxFun) -> DirAlias
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IxFun -> FlatSlice (TPrimExp Int64 VName) -> IxFun
forall num.
(Eq num, IntegralExp num) =>
IxFun num -> FlatSlice num -> IxFun num
`IxFun.flatSlice` TPrimExp Int64 VName
-> [FlatDimIndex (TPrimExp Int64 VName)]
-> FlatSlice (TPrimExp Int64 VName)
forall d. d -> [FlatDimIndex d] -> FlatSlice d
FlatSlice (SubExp -> TPrimExp Int64 VName
pe64 SubExp
offset) ((FlatDimIndex SubExp -> FlatDimIndex (TPrimExp Int64 VName))
-> [FlatDimIndex SubExp] -> [FlatDimIndex (TPrimExp Int64 VName)]
forall a b. (a -> b) -> [a] -> [b]
map ((SubExp -> TPrimExp Int64 VName)
-> FlatDimIndex SubExp -> FlatDimIndex (TPrimExp Int64 VName)
forall a b. (a -> b) -> FlatDimIndex a -> FlatDimIndex b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64) [FlatDimIndex SubExp]
idxs))
)
getDirAliasFromExp (BasicOp (FlatUpdate VName
x FlatSlice SubExp
_ VName
_)) = (VName, DirAlias) -> Maybe (VName, DirAlias)
forall a. a -> Maybe a
Just (VName
x, DirAlias
forall a. a -> Maybe a
Just)
getDirAliasFromExp Exp (Aliases rep)
_ = Maybe (VName, DirAlias)
forall a. Maybe a
Nothing
getInvAliasFromExp :: Exp (Aliases rep) -> InvAlias
getInvAliasFromExp :: forall rep. Exp (Aliases rep) -> InvAlias
getInvAliasFromExp (BasicOp (SubExp (Var VName
_))) = (IxFun -> IxFun) -> InvAlias
forall a. a -> Maybe a
Just IxFun -> IxFun
forall a. a -> a
id
getInvAliasFromExp (BasicOp (Opaque OpaqueOp
_ (Var VName
_))) = (IxFun -> IxFun) -> InvAlias
forall a. a -> Maybe a
Just IxFun -> IxFun
forall a. a -> a
id
getInvAliasFromExp (BasicOp Update {}) = (IxFun -> IxFun) -> InvAlias
forall a. a -> Maybe a
Just IxFun -> IxFun
forall a. a -> a
id
getInvAliasFromExp (BasicOp (Rearrange [Int]
perm VName
_)) =
(IxFun -> IxFun) -> InvAlias
forall a. a -> Maybe a
Just (IxFun -> [Int] -> IxFun
forall num. IntegralExp num => IxFun num -> [Int] -> IxFun num
`IxFun.permute` [Int] -> [Int]
rearrangeInverse [Int]
perm)
getInvAliasFromExp Exp (Aliases rep)
_ = InvAlias
forall a. Maybe a
Nothing
class TopDownHelper inner where
innerNonNegatives :: [VName] -> inner -> Names
innerKnownLessThan :: inner -> [(VName, PrimExp VName)]
scopeHelper :: inner -> Scope rep
instance TopDownHelper (SegOp lvl rep) where
innerNonNegatives :: [VName] -> SegOp lvl rep -> Names
innerNonNegatives [VName]
_ SegOp lvl rep
op =
((VName, SubExp) -> Names) -> [(VName, SubExp)] -> Names
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (VName -> Names
oneName (VName -> Names)
-> ((VName, SubExp) -> VName) -> (VName, SubExp) -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, SubExp) -> VName
forall a b. (a, b) -> a
fst) ([(VName, SubExp)] -> Names) -> [(VName, SubExp)] -> Names
forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace (SegSpace -> [(VName, SubExp)]) -> SegSpace -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$ SegOp lvl rep -> SegSpace
forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp lvl rep
op
innerKnownLessThan :: SegOp lvl rep -> [(VName, PrimExp VName)]
innerKnownLessThan SegOp lvl rep
op =
((VName, SubExp) -> (VName, PrimExp VName))
-> [(VName, SubExp)] -> [(VName, PrimExp VName)]
forall a b. (a -> b) -> [a] -> [b]
map ((SubExp -> PrimExp VName)
-> (VName, SubExp) -> (VName, PrimExp VName)
forall a b. (a -> b) -> (VName, a) -> (VName, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((SubExp -> PrimExp VName)
-> (VName, SubExp) -> (VName, PrimExp VName))
-> (SubExp -> PrimExp VName)
-> (VName, SubExp)
-> (VName, PrimExp VName)
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> PrimExp VName
primExpFromSubExp (PrimType -> SubExp -> PrimExp VName)
-> PrimType -> SubExp -> PrimExp VName
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
Int64) ([(VName, SubExp)] -> [(VName, PrimExp VName)])
-> [(VName, SubExp)] -> [(VName, PrimExp VName)]
forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace (SegSpace -> [(VName, SubExp)]) -> SegSpace -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$ SegOp lvl rep -> SegSpace
forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp lvl rep
op
scopeHelper :: forall rep. SegOp lvl rep -> Scope rep
scopeHelper SegOp lvl rep
op = SegSpace -> Scope rep
forall rep. SegSpace -> Scope rep
scopeOfSegSpace (SegSpace -> Scope rep) -> SegSpace -> Scope rep
forall a b. (a -> b) -> a -> b
$ SegOp lvl rep -> SegSpace
forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp lvl rep
op
instance TopDownHelper (HostOp NoOp (Aliases GPUMem)) where
innerNonNegatives :: [VName] -> HostOp NoOp (Aliases GPUMem) -> Names
innerNonNegatives [VName]
vs (SegOp SegOp SegLevel (Aliases GPUMem)
op) = [VName] -> SegOp SegLevel (Aliases GPUMem) -> Names
forall inner. TopDownHelper inner => [VName] -> inner -> Names
innerNonNegatives [VName]
vs SegOp SegLevel (Aliases GPUMem)
op
innerNonNegatives [VName
vname] (SizeOp (GetSize Name
_ SizeClass
_)) = VName -> Names
oneName VName
vname
innerNonNegatives [VName
vname] (SizeOp (GetSizeMax SizeClass
_)) = VName -> Names
oneName VName
vname
innerNonNegatives [VName]
_ HostOp NoOp (Aliases GPUMem)
_ = Names
forall a. Monoid a => a
mempty
innerKnownLessThan :: HostOp NoOp (Aliases GPUMem) -> [(VName, PrimExp VName)]
innerKnownLessThan (SegOp SegOp SegLevel (Aliases GPUMem)
op) = SegOp SegLevel (Aliases GPUMem) -> [(VName, PrimExp VName)]
forall inner.
TopDownHelper inner =>
inner -> [(VName, PrimExp VName)]
innerKnownLessThan SegOp SegLevel (Aliases GPUMem)
op
innerKnownLessThan HostOp NoOp (Aliases GPUMem)
_ = [(VName, PrimExp VName)]
forall a. Monoid a => a
mempty
scopeHelper :: forall rep. HostOp NoOp (Aliases GPUMem) -> Scope rep
scopeHelper (SegOp SegOp SegLevel (Aliases GPUMem)
op) = SegOp SegLevel (Aliases GPUMem) -> Scope rep
forall rep. SegOp SegLevel (Aliases GPUMem) -> Scope rep
forall inner rep. TopDownHelper inner => inner -> Scope rep
scopeHelper SegOp SegLevel (Aliases GPUMem)
op
scopeHelper HostOp NoOp (Aliases GPUMem)
_ = Scope rep
forall a. Monoid a => a
mempty
instance (TopDownHelper (inner (Aliases MCMem))) => TopDownHelper (MC.MCOp inner (Aliases MCMem)) where
innerNonNegatives :: [VName] -> MCOp inner (Aliases MCMem) -> Names
innerNonNegatives [VName]
vs (ParOp Maybe (SegOp () (Aliases MCMem))
par_op SegOp () (Aliases MCMem)
op) =
Names
-> (SegOp () (Aliases MCMem) -> Names)
-> Maybe (SegOp () (Aliases MCMem))
-> Names
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Names
forall a. Monoid a => a
mempty ([VName] -> SegOp () (Aliases MCMem) -> Names
forall inner. TopDownHelper inner => [VName] -> inner -> Names
innerNonNegatives [VName]
vs) Maybe (SegOp () (Aliases MCMem))
par_op
Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [VName] -> SegOp () (Aliases MCMem) -> Names
forall inner. TopDownHelper inner => [VName] -> inner -> Names
innerNonNegatives [VName]
vs SegOp () (Aliases MCMem)
op
innerNonNegatives [VName]
vs (MC.OtherOp inner (Aliases MCMem)
op) =
[VName] -> inner (Aliases MCMem) -> Names
forall inner. TopDownHelper inner => [VName] -> inner -> Names
innerNonNegatives [VName]
vs inner (Aliases MCMem)
op
innerKnownLessThan :: MCOp inner (Aliases MCMem) -> [(VName, PrimExp VName)]
innerKnownLessThan (ParOp Maybe (SegOp () (Aliases MCMem))
par_op SegOp () (Aliases MCMem)
op) =
[(VName, PrimExp VName)]
-> (SegOp () (Aliases MCMem) -> [(VName, PrimExp VName)])
-> Maybe (SegOp () (Aliases MCMem))
-> [(VName, PrimExp VName)]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [(VName, PrimExp VName)]
forall a. Monoid a => a
mempty SegOp () (Aliases MCMem) -> [(VName, PrimExp VName)]
forall inner.
TopDownHelper inner =>
inner -> [(VName, PrimExp VName)]
innerKnownLessThan Maybe (SegOp () (Aliases MCMem))
par_op [(VName, PrimExp VName)]
-> [(VName, PrimExp VName)] -> [(VName, PrimExp VName)]
forall a. Semigroup a => a -> a -> a
<> SegOp () (Aliases MCMem) -> [(VName, PrimExp VName)]
forall inner.
TopDownHelper inner =>
inner -> [(VName, PrimExp VName)]
innerKnownLessThan SegOp () (Aliases MCMem)
op
innerKnownLessThan (MC.OtherOp inner (Aliases MCMem)
op) =
inner (Aliases MCMem) -> [(VName, PrimExp VName)]
forall inner.
TopDownHelper inner =>
inner -> [(VName, PrimExp VName)]
innerKnownLessThan inner (Aliases MCMem)
op
scopeHelper :: forall rep. MCOp inner (Aliases MCMem) -> Scope rep
scopeHelper (ParOp Maybe (SegOp () (Aliases MCMem))
par_op SegOp () (Aliases MCMem)
op) =
Scope rep
-> (SegOp () (Aliases MCMem) -> Scope rep)
-> Maybe (SegOp () (Aliases MCMem))
-> Scope rep
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Scope rep
forall a. Monoid a => a
mempty SegOp () (Aliases MCMem) -> Scope rep
forall rep. SegOp () (Aliases MCMem) -> Scope rep
forall inner rep. TopDownHelper inner => inner -> Scope rep
scopeHelper Maybe (SegOp () (Aliases MCMem))
par_op Scope rep -> Scope rep -> Scope rep
forall a. Semigroup a => a -> a -> a
<> SegOp () (Aliases MCMem) -> Scope rep
forall rep. SegOp () (Aliases MCMem) -> Scope rep
forall inner rep. TopDownHelper inner => inner -> Scope rep
scopeHelper SegOp () (Aliases MCMem)
op
scopeHelper MC.OtherOp {} = Scope rep
forall a. Monoid a => a
mempty
instance TopDownHelper (NoOp rep) where
innerNonNegatives :: [VName] -> NoOp rep -> Names
innerNonNegatives [VName]
_ NoOp rep
NoOp = Names
forall a. Monoid a => a
mempty
innerKnownLessThan :: NoOp rep -> [(VName, PrimExp VName)]
innerKnownLessThan NoOp rep
NoOp = [(VName, PrimExp VName)]
forall a. Monoid a => a
mempty
scopeHelper :: forall rep. NoOp rep -> Scope rep
scopeHelper NoOp rep
NoOp = Scope rep
forall a. Monoid a => a
mempty
updateTopdownEnv ::
(ASTRep rep, Op rep ~ MemOp inner rep, TopDownHelper (inner (Aliases rep))) =>
TopdownEnv rep ->
Stm (Aliases rep) ->
TopdownEnv rep
updateTopdownEnv :: forall rep (inner :: * -> *).
(ASTRep rep, Op rep ~ MemOp inner rep,
TopDownHelper (inner (Aliases rep))) =>
TopdownEnv rep -> Stm (Aliases rep) -> TopdownEnv rep
updateTopdownEnv TopdownEnv rep
env stm :: Stm (Aliases rep)
stm@(Let (Pat [PatElem (LetDec (Aliases rep))
pe]) StmAux (ExpDec (Aliases rep))
_ (Op (Alloc (Var VName
vname) Space
sp))) =
TopdownEnv rep
env
{ alloc :: AllocTab
alloc = VName -> Space -> AllocTab -> AllocTab
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (PatElem (VarAliases, LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (VarAliases, LetDec rep)
PatElem (LetDec (Aliases rep))
pe) Space
sp (AllocTab -> AllocTab) -> AllocTab -> AllocTab
forall a b. (a -> b) -> a -> b
$ TopdownEnv rep -> AllocTab
forall rep. TopdownEnv rep -> AllocTab
alloc TopdownEnv rep
env,
scope :: ScopeTab rep
scope = TopdownEnv rep -> ScopeTab rep
forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
env ScopeTab rep -> ScopeTab rep -> ScopeTab rep
forall a. Semigroup a => a -> a -> a
<> Stm (Aliases rep) -> ScopeTab rep
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stm (Aliases rep)
stm,
nonNegatives :: Names
nonNegatives = TopdownEnv rep -> Names
forall rep. TopdownEnv rep -> Names
nonNegatives TopdownEnv rep
env Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> VName -> Names
oneName VName
vname
}
updateTopdownEnv TopdownEnv rep
env stm :: Stm (Aliases rep)
stm@(Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
_ (Op (Inner inner (Aliases rep)
inner))) =
TopdownEnv rep
env
{ scope :: ScopeTab rep
scope = TopdownEnv rep -> ScopeTab rep
forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
env ScopeTab rep -> ScopeTab rep -> ScopeTab rep
forall a. Semigroup a => a -> a -> a
<> Stm (Aliases rep) -> ScopeTab rep
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stm (Aliases rep)
stm ScopeTab rep -> ScopeTab rep -> ScopeTab rep
forall a. Semigroup a => a -> a -> a
<> inner (Aliases rep) -> ScopeTab rep
forall rep. inner (Aliases rep) -> Scope rep
forall inner rep. TopDownHelper inner => inner -> Scope rep
scopeHelper inner (Aliases rep)
inner,
nonNegatives :: Names
nonNegatives = TopdownEnv rep -> Names
forall rep. TopdownEnv rep -> Names
nonNegatives TopdownEnv rep
env Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [VName] -> inner (Aliases rep) -> Names
forall inner. TopDownHelper inner => [VName] -> inner -> Names
innerNonNegatives (Pat (VarAliases, LetDec rep) -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat (VarAliases, LetDec rep)
Pat (LetDec (Aliases rep))
pat) inner (Aliases rep)
inner,
knownLessThan :: [(VName, PrimExp VName)]
knownLessThan = TopdownEnv rep -> [(VName, PrimExp VName)]
forall rep. TopdownEnv rep -> [(VName, PrimExp VName)]
knownLessThan TopdownEnv rep
env [(VName, PrimExp VName)]
-> [(VName, PrimExp VName)] -> [(VName, PrimExp VName)]
forall a. Semigroup a => a -> a -> a
<> inner (Aliases rep) -> [(VName, PrimExp VName)]
forall inner.
TopDownHelper inner =>
inner -> [(VName, PrimExp VName)]
innerKnownLessThan inner (Aliases rep)
inner
}
updateTopdownEnv TopdownEnv rep
env stm :: Stm (Aliases rep)
stm@(Let (Pat [PatElem (LetDec (Aliases rep))]
_) StmAux (ExpDec (Aliases rep))
_ (BasicOp (Assert SubExp
se ErrorMsg SubExp
_ (SrcLoc, [SrcLoc])
_))) =
TopdownEnv rep
env
{ scope :: ScopeTab rep
scope = TopdownEnv rep -> ScopeTab rep
forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
env ScopeTab rep -> ScopeTab rep -> ScopeTab rep
forall a. Semigroup a => a -> a -> a
<> Stm (Aliases rep) -> ScopeTab rep
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stm (Aliases rep)
stm,
td_asserts :: [SubExp]
td_asserts = SubExp
se SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: TopdownEnv rep -> [SubExp]
forall rep. TopdownEnv rep -> [SubExp]
td_asserts TopdownEnv rep
env
}
updateTopdownEnv TopdownEnv rep
env stm :: Stm (Aliases rep)
stm@(Let (Pat [PatElem (LetDec (Aliases rep))
pe]) StmAux (ExpDec (Aliases rep))
_ Exp (Aliases rep)
e)
| Just (VName
x, DirAlias
ixfn) <- Exp (Aliases rep) -> Maybe (VName, DirAlias)
forall rep. Exp (Aliases rep) -> Maybe (VName, DirAlias)
getDirAliasFromExp Exp (Aliases rep)
e =
let ixfn_inv :: InvAlias
ixfn_inv = Exp (Aliases rep) -> InvAlias
forall rep. Exp (Aliases rep) -> InvAlias
getInvAliasFromExp Exp (Aliases rep)
e
in TopdownEnv rep
env
{ v_alias :: VarAliasTab
v_alias = VName -> (VName, DirAlias, InvAlias) -> VarAliasTab -> VarAliasTab
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (PatElem (VarAliases, LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (VarAliases, LetDec rep)
PatElem (LetDec (Aliases rep))
pe) (VName
x, DirAlias
ixfn, InvAlias
ixfn_inv) (TopdownEnv rep -> VarAliasTab
forall rep. TopdownEnv rep -> VarAliasTab
v_alias TopdownEnv rep
env),
scope :: ScopeTab rep
scope = TopdownEnv rep -> ScopeTab rep
forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
env ScopeTab rep -> ScopeTab rep -> ScopeTab rep
forall a. Semigroup a => a -> a -> a
<> Stm (Aliases rep) -> ScopeTab rep
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stm (Aliases rep)
stm,
nonNegatives :: Names
nonNegatives = TopdownEnv rep -> Names
forall rep. TopdownEnv rep -> Names
nonNegatives TopdownEnv rep
env Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Pat (VarAliases, LetDec rep) -> Names
forall rep. Typed rep => Pat rep -> Names
nonNegativesInPat (Stm (Aliases rep) -> Pat (LetDec (Aliases rep))
forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm (Aliases rep)
stm)
}
updateTopdownEnv TopdownEnv rep
env Stm (Aliases rep)
stm =
TopdownEnv rep
env
{ scope :: ScopeTab rep
scope = TopdownEnv rep -> ScopeTab rep
forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
env ScopeTab rep -> ScopeTab rep -> ScopeTab rep
forall a. Semigroup a => a -> a -> a
<> Stm (Aliases rep) -> ScopeTab rep
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stm (Aliases rep)
stm,
nonNegatives :: Names
nonNegatives = TopdownEnv rep -> Names
forall rep. TopdownEnv rep -> Names
nonNegatives TopdownEnv rep
env Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Pat (VarAliases, LetDec rep) -> Names
forall rep. Typed rep => Pat rep -> Names
nonNegativesInPat (Stm (Aliases rep) -> Pat (LetDec (Aliases rep))
forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm (Aliases rep)
stm)
}
nonNegativesInPat :: (Typed rep) => Pat rep -> Names
nonNegativesInPat :: forall rep. Typed rep => Pat rep -> Names
nonNegativesInPat (Pat [PatElem rep]
elems) =
(PatElem rep -> Names) -> [PatElem rep] -> Names
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ([VName] -> Names
namesFromList ([VName] -> Names)
-> (PatElem rep -> [VName]) -> PatElem rep -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp -> Maybe VName) -> [SubExp] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe SubExp -> Maybe VName
subExpVar ([SubExp] -> [VName])
-> (PatElem rep -> [SubExp]) -> PatElem rep -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
-> (PatElem rep -> TypeBase Shape NoUniqueness)
-> PatElem rep
-> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatElem rep -> TypeBase Shape NoUniqueness
forall t. Typed t => t -> TypeBase Shape NoUniqueness
typeOf) [PatElem rep]
elems
updateTopdownEnvLoop :: TopdownEnv rep -> [(FParam rep, SubExp)] -> LoopForm -> TopdownEnv rep
updateTopdownEnvLoop :: forall rep.
TopdownEnv rep
-> [(FParam rep, SubExp)] -> LoopForm -> TopdownEnv rep
updateTopdownEnvLoop TopdownEnv rep
td_env [(FParam rep, SubExp)]
arginis LoopForm
lform =
let scopetab :: ScopeTab rep
scopetab =
TopdownEnv rep -> ScopeTab rep
forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env
ScopeTab rep -> ScopeTab rep -> ScopeTab rep
forall a. Semigroup a => a -> a -> a
<> [FParam rep] -> ScopeTab rep
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams (((FParam rep, SubExp) -> FParam rep)
-> [(FParam rep, SubExp)] -> [FParam rep]
forall a b. (a -> b) -> [a] -> [b]
map (FParam rep, SubExp) -> FParam rep
forall a b. (a, b) -> a
fst [(FParam rep, SubExp)]
arginis)
ScopeTab rep -> ScopeTab rep -> ScopeTab rep
forall a. Semigroup a => a -> a -> a
<> LoopForm -> ScopeTab rep
forall rep. LoopForm -> Scope rep
scopeOfLoopForm LoopForm
lform
non_negatives :: Names
non_negatives =
TopdownEnv rep -> Names
forall rep. TopdownEnv rep -> Names
nonNegatives TopdownEnv rep
td_env Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> case LoopForm
lform of
ForLoop VName
v IntType
_ SubExp
_ -> VName -> Names
oneName VName
v
LoopForm
_ -> Names
forall a. Monoid a => a
mempty
less_than :: [(VName, PrimExp VName)]
less_than =
case LoopForm
lform of
ForLoop VName
v IntType
_ SubExp
b -> [(VName
v, PrimType -> SubExp -> PrimExp VName
primExpFromSubExp (IntType -> PrimType
IntType IntType
Int64) SubExp
b)]
LoopForm
_ -> [(VName, PrimExp VName)]
forall a. Monoid a => a
mempty
in TopdownEnv rep
td_env
{ scope :: ScopeTab rep
scope = ScopeTab rep
scopetab,
nonNegatives :: Names
nonNegatives = Names
non_negatives,
knownLessThan :: [(VName, PrimExp VName)]
knownLessThan = [(VName, PrimExp VName)]
less_than [(VName, PrimExp VName)]
-> [(VName, PrimExp VName)] -> [(VName, PrimExp VName)]
forall a. Semigroup a => a -> a -> a
<> TopdownEnv rep -> [(VName, PrimExp VName)]
forall rep. TopdownEnv rep -> [(VName, PrimExp VName)]
knownLessThan TopdownEnv rep
td_env
}
getDirAliasedIxfn :: (HasMemBlock (Aliases rep)) => TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, IxFun)
getDirAliasedIxfn :: forall rep.
HasMemBlock (Aliases rep) =>
TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, IxFun)
getDirAliasedIxfn TopdownEnv rep
td_env CoalsTab
coals_tab VName
x =
case VName -> Scope (Aliases rep) -> Maybe ArrayMemBound
forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
x (TopdownEnv rep -> Scope (Aliases rep)
forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) of
Just (MemBlock PrimType
_ Shape
_ VName
m_x IxFun
orig_ixfun) ->
case VName -> CoalsTab -> Maybe CoalsEntry
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_x CoalsTab
coals_tab of
Just CoalsEntry
coal_etry -> do
(Coalesced CoalescedKind
_ (MemBlock PrimType
_ Shape
_ VName
m IxFun
ixf) FreeVarSubsts
_) <- VarAliasTab -> Map VName Coalesced -> VName -> Maybe Coalesced
walkAliasTab (TopdownEnv rep -> VarAliasTab
forall rep. TopdownEnv rep -> VarAliasTab
v_alias TopdownEnv rep
td_env) (CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
coal_etry) VName
x
(VName, VName, IxFun) -> Maybe (VName, VName, IxFun)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
m_x, VName
m, IxFun
ixf)
Maybe CoalsEntry
Nothing ->
(VName, VName, IxFun) -> Maybe (VName, VName, IxFun)
forall a. a -> Maybe a
Just (VName
m_x, VName
m_x, IxFun
orig_ixfun)
Maybe ArrayMemBound
Nothing -> Maybe (VName, VName, IxFun)
forall a. Maybe a
Nothing
getDirAliasedIxfn' :: (HasMemBlock (Aliases rep)) => TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, IxFun)
getDirAliasedIxfn' :: forall rep.
HasMemBlock (Aliases rep) =>
TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, IxFun)
getDirAliasedIxfn' TopdownEnv rep
td_env CoalsTab
coals_tab VName
x =
case VName -> Scope (Aliases rep) -> Maybe ArrayMemBound
forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
x (TopdownEnv rep -> Scope (Aliases rep)
forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) of
Just (MemBlock PrimType
_ Shape
_ VName
m_x IxFun
_) ->
case VName -> CoalsTab -> Maybe CoalsEntry
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
m_x CoalsTab
coals_tab of
Just CoalsEntry
coal_etry -> do
(Coalesced CoalescedKind
_ (MemBlock PrimType
_ Shape
_ VName
m IxFun
ixf) FreeVarSubsts
_) <- VarAliasTab -> Map VName Coalesced -> VName -> Maybe Coalesced
walkAliasTab (TopdownEnv rep -> VarAliasTab
forall rep. TopdownEnv rep -> VarAliasTab
v_alias TopdownEnv rep
td_env) (CoalsEntry -> Map VName Coalesced
vartab CoalsEntry
coal_etry) VName
x
(VName, VName, IxFun) -> Maybe (VName, VName, IxFun)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
m_x, VName
m, IxFun
ixf)
Maybe CoalsEntry
Nothing ->
Maybe (VName, VName, IxFun)
forall a. Maybe a
Nothing
Maybe ArrayMemBound
Nothing -> Maybe (VName, VName, IxFun)
forall a. Maybe a
Nothing
walkAliasTab ::
VarAliasTab ->
M.Map VName Coalesced ->
VName ->
Maybe Coalesced
walkAliasTab :: VarAliasTab -> Map VName Coalesced -> VName -> Maybe Coalesced
walkAliasTab VarAliasTab
_ Map VName Coalesced
vtab VName
x
| Just Coalesced
c <- VName -> Map VName Coalesced -> Maybe Coalesced
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x Map VName Coalesced
vtab =
Coalesced -> Maybe Coalesced
forall a. a -> Maybe a
Just Coalesced
c
walkAliasTab VarAliasTab
alias_tab Map VName Coalesced
vtab VName
x
| Just (VName
x0, DirAlias
alias0, InvAlias
_) <- VName -> VarAliasTab -> Maybe (VName, DirAlias, InvAlias)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x VarAliasTab
alias_tab = do
Coalesced CoalescedKind
knd (MemBlock PrimType
pt Shape
shp VName
vname IxFun
ixf) FreeVarSubsts
substs <- VarAliasTab -> Map VName Coalesced -> VName -> Maybe Coalesced
walkAliasTab VarAliasTab
alias_tab Map VName Coalesced
vtab VName
x0
IxFun
ixf' <- DirAlias
alias0 IxFun
ixf
Coalesced -> Maybe Coalesced
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Coalesced -> Maybe Coalesced) -> Coalesced -> Maybe Coalesced
forall a b. (a -> b) -> a -> b
$ CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced CoalescedKind
knd (PrimType -> Shape -> VName -> IxFun -> ArrayMemBound
MemBlock PrimType
pt Shape
shp VName
vname IxFun
ixf') FreeVarSubsts
substs
walkAliasTab VarAliasTab
_ Map VName Coalesced
_ VName
_ = Maybe Coalesced
forall a. Maybe a
Nothing
addInvAliasesVarTab ::
(HasMemBlock (Aliases rep)) =>
TopdownEnv rep ->
M.Map VName Coalesced ->
VName ->
Maybe (M.Map VName Coalesced)
addInvAliasesVarTab :: 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
x
| Just (Coalesced CoalescedKind
_ (MemBlock PrimType
_ Shape
_ VName
m_y IxFun
x_ixfun) FreeVarSubsts
fv_subs) <- VName -> Map VName Coalesced -> Maybe Coalesced
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x Map VName Coalesced
vtab =
case VName -> VarAliasTab -> Maybe (VName, DirAlias, InvAlias)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x (TopdownEnv rep -> VarAliasTab
forall rep. TopdownEnv rep -> VarAliasTab
v_alias TopdownEnv rep
td_env) of
Maybe (VName, DirAlias, InvAlias)
Nothing -> Map VName Coalesced -> Maybe (Map VName Coalesced)
forall a. a -> Maybe a
Just Map VName Coalesced
vtab
Just (VName
_, DirAlias
_, InvAlias
Nothing) -> Maybe (Map VName Coalesced)
forall a. Maybe a
Nothing
Just (VName
x0, DirAlias
_, Just IxFun -> IxFun
inv_alias0) ->
let x_ixfn0 :: IxFun
x_ixfn0 = IxFun -> IxFun
inv_alias0 IxFun
x_ixfun
in case VName -> Scope (Aliases rep) -> Maybe ArrayMemBound
forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
getScopeMemInfo VName
x0 (TopdownEnv rep -> Scope (Aliases rep)
forall rep. TopdownEnv rep -> ScopeTab rep
scope TopdownEnv rep
td_env) of
Maybe ArrayMemBound
Nothing -> [Char] -> Maybe (Map VName Coalesced)
forall a. HasCallStack => [Char] -> a
error [Char]
"impossible"
Just (MemBlock PrimType
ptp Shape
shp VName
_ IxFun
_) ->
let coal :: Coalesced
coal = CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
Coalesced CoalescedKind
TransitiveCoal (PrimType -> Shape -> VName -> IxFun -> ArrayMemBound
MemBlock PrimType
ptp Shape
shp VName
m_y IxFun
x_ixfn0) FreeVarSubsts
fv_subs
vartab' :: Map VName Coalesced
vartab' = VName -> Coalesced -> Map VName Coalesced -> Map VName Coalesced
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
x0 Coalesced
coal Map VName Coalesced
vtab
in TopdownEnv rep
-> Map VName Coalesced -> VName -> Maybe (Map VName Coalesced)
forall rep.
HasMemBlock (Aliases rep) =>
TopdownEnv rep
-> Map VName Coalesced -> VName -> Maybe (Map VName Coalesced)
addInvAliasesVarTab TopdownEnv rep
td_env Map VName Coalesced
vartab' VName
x0
addInvAliasesVarTab TopdownEnv rep
_ Map VName Coalesced
_ VName
_ = Maybe (Map VName Coalesced)
forall a. Maybe a
Nothing
areAliased :: TopdownEnv rep -> VName -> VName -> Bool
areAliased :: forall rep. TopdownEnv rep -> VName -> VName -> Bool
areAliased TopdownEnv rep
_ VName
m_x VName
m_y =
VName
m_x VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
m_y
areAnyAliased :: TopdownEnv rep -> VName -> [VName] -> Bool
areAnyAliased :: forall rep. TopdownEnv rep -> VName -> [VName] -> Bool
areAnyAliased TopdownEnv rep
td_env VName
m_x =
(VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (TopdownEnv rep -> VName -> VName -> Bool
forall rep. TopdownEnv rep -> VName -> VName -> Bool
areAliased TopdownEnv rep
td_env VName
m_x)