{-# LANGUAGE TypeFamilies #-}
module Futhark.Analysis.SymbolTable
( SymbolTable (bindings, loopDepth, availableAtClosestLoop, simplifyMemory),
empty,
fromScope,
toScope,
Entry,
deepen,
entryAccInput,
entryDepth,
entryLetBoundDec,
entryIsSize,
entryStm,
entryFParam,
entryLParam,
elem,
lookup,
lookupStm,
lookupExp,
lookupBasicOp,
lookupType,
lookupSubExp,
lookupAliases,
lookupLoopVar,
lookupLoopParam,
aliases,
available,
subExpAvailable,
consume,
index,
index',
Indexed (..),
indexedAddCerts,
IndexOp (..),
insertStm,
insertStms,
insertFParams,
insertLParam,
insertLoopVar,
insertLoopMerge,
hideCertified,
noteAccTokens,
)
where
import Control.Arrow ((&&&))
import Control.Monad
import Data.List (elemIndex, foldl')
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Ord
import Futhark.Analysis.PrimExp.Convert
import Futhark.IR hiding (FParam, lookupType)
import Futhark.IR qualified as AST
import Futhark.IR.Prop.Aliases qualified as Aliases
import Prelude hiding (elem, lookup)
data SymbolTable rep = SymbolTable
{ forall rep. SymbolTable rep -> Int
loopDepth :: Int,
forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings :: M.Map VName (Entry rep),
forall rep. SymbolTable rep -> Names
availableAtClosestLoop :: Names,
forall rep. SymbolTable rep -> Bool
simplifyMemory :: Bool
}
instance Semigroup (SymbolTable rep) where
SymbolTable rep
table1 <> :: SymbolTable rep -> SymbolTable rep -> SymbolTable rep
<> SymbolTable rep
table2 =
SymbolTable
{ loopDepth :: Int
loopDepth = forall a. Ord a => a -> a -> a
max (forall rep. SymbolTable rep -> Int
loopDepth SymbolTable rep
table1) (forall rep. SymbolTable rep -> Int
loopDepth SymbolTable rep
table2),
bindings :: Map VName (Entry rep)
bindings = forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
table1 forall a. Semigroup a => a -> a -> a
<> forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
table2,
availableAtClosestLoop :: Names
availableAtClosestLoop =
forall rep. SymbolTable rep -> Names
availableAtClosestLoop SymbolTable rep
table1
forall a. Semigroup a => a -> a -> a
<> forall rep. SymbolTable rep -> Names
availableAtClosestLoop SymbolTable rep
table2,
simplifyMemory :: Bool
simplifyMemory = forall rep. SymbolTable rep -> Bool
simplifyMemory SymbolTable rep
table1 Bool -> Bool -> Bool
|| forall rep. SymbolTable rep -> Bool
simplifyMemory SymbolTable rep
table2
}
instance Monoid (SymbolTable rep) where
mempty :: SymbolTable rep
mempty = forall rep. SymbolTable rep
empty
empty :: SymbolTable rep
empty :: forall rep. SymbolTable rep
empty = forall rep.
Int -> Map VName (Entry rep) -> Names -> Bool -> SymbolTable rep
SymbolTable Int
0 forall k a. Map k a
M.empty forall a. Monoid a => a
mempty Bool
False
fromScope :: ASTRep rep => Scope rep -> SymbolTable rep
fromScope :: forall rep. ASTRep rep => Scope rep -> SymbolTable rep
fromScope = forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey' forall {rep}.
ASTRep rep =>
SymbolTable rep -> VName -> NameInfo rep -> SymbolTable rep
insertFreeVar' forall rep. SymbolTable rep
empty
where
insertFreeVar' :: SymbolTable rep -> VName -> NameInfo rep -> SymbolTable rep
insertFreeVar' SymbolTable rep
m VName
k NameInfo rep
dec = forall rep.
ASTRep rep =>
VName -> NameInfo rep -> SymbolTable rep -> SymbolTable rep
insertFreeVar VName
k NameInfo rep
dec SymbolTable rep
m
toScope :: SymbolTable rep -> Scope rep
toScope :: forall rep. SymbolTable rep -> Scope rep
toScope = forall a b k. (a -> b) -> Map k a -> Map k b
M.map forall rep. Entry rep -> NameInfo rep
entryInfo forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings
deepen :: SymbolTable rep -> SymbolTable rep
deepen :: forall rep. SymbolTable rep -> SymbolTable rep
deepen SymbolTable rep
vtable =
SymbolTable rep
vtable
{ loopDepth :: Int
loopDepth = forall rep. SymbolTable rep -> Int
loopDepth SymbolTable rep
vtable forall a. Num a => a -> a -> a
+ Int
1,
availableAtClosestLoop :: Names
availableAtClosestLoop = [VName] -> Names
namesFromList forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [k]
M.keys forall a b. (a -> b) -> a -> b
$ forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable
}
data Indexed
=
Indexed Certs (PrimExp VName)
|
IndexedArray Certs VName [TPrimExp Int64 VName]
indexedAddCerts :: Certs -> Indexed -> Indexed
indexedAddCerts :: Certs -> Indexed -> Indexed
indexedAddCerts Certs
cs1 (Indexed Certs
cs2 PrimExp VName
v) = Certs -> PrimExp VName -> Indexed
Indexed (Certs
cs1 forall a. Semigroup a => a -> a -> a
<> Certs
cs2) PrimExp VName
v
indexedAddCerts Certs
cs1 (IndexedArray Certs
cs2 VName
arr [TPrimExp Int64 VName]
v) = Certs -> VName -> [TPrimExp Int64 VName] -> Indexed
IndexedArray (Certs
cs1 forall a. Semigroup a => a -> a -> a
<> Certs
cs2) VName
arr [TPrimExp Int64 VName]
v
instance FreeIn Indexed where
freeIn' :: Indexed -> FV
freeIn' (Indexed Certs
cs PrimExp VName
v) = forall a. FreeIn a => a -> FV
freeIn' Certs
cs forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' PrimExp VName
v
freeIn' (IndexedArray Certs
cs VName
arr [TPrimExp Int64 VName]
v) = forall a. FreeIn a => a -> FV
freeIn' Certs
cs forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' VName
arr forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' [TPrimExp Int64 VName]
v
type IndexArray = [TPrimExp Int64 VName] -> Maybe Indexed
data Entry rep = Entry
{
forall rep. Entry rep -> Bool
entryConsumed :: Bool,
forall rep. Entry rep -> Int
entryDepth :: Int,
forall rep. Entry rep -> Bool
entryIsSize :: Bool,
forall rep. Entry rep -> Maybe (WithAccInput rep)
entryAccInput :: Maybe (WithAccInput rep),
forall rep. Entry rep -> EntryType rep
entryType :: EntryType rep
}
data EntryType rep
= LoopVar (LoopVarEntry rep)
| LetBound (LetBoundEntry rep)
| FParam (FParamEntry rep)
| LParam (LParamEntry rep)
| FreeVar (FreeVarEntry rep)
data LoopVarEntry rep = LoopVarEntry
{ forall {k} (rep :: k). LoopVarEntry rep -> IntType
loopVarType :: IntType,
forall {k} (rep :: k). LoopVarEntry rep -> SubExp
loopVarBound :: SubExp
}
data LetBoundEntry rep = LetBoundEntry
{ forall rep. LetBoundEntry rep -> LetDec rep
letBoundDec :: LetDec rep,
forall rep. LetBoundEntry rep -> Names
letBoundAliases :: Names,
forall rep. LetBoundEntry rep -> Stm rep
letBoundStm :: Stm rep,
forall rep. LetBoundEntry rep -> Int -> IndexArray
letBoundIndex :: Int -> IndexArray
}
data FParamEntry rep = FParamEntry
{ forall rep. FParamEntry rep -> FParamInfo rep
fparamDec :: FParamInfo rep,
forall rep. FParamEntry rep -> Names
fparamAliases :: Names,
forall rep. FParamEntry rep -> Maybe (SubExp, SubExp)
fparamMerge :: Maybe (SubExp, SubExp)
}
data LParamEntry rep = LParamEntry
{ forall rep. LParamEntry rep -> LParamInfo rep
lparamDec :: LParamInfo rep,
forall rep. LParamEntry rep -> Names
lparamAliases :: Names,
forall rep. LParamEntry rep -> IndexArray
lparamIndex :: IndexArray
}
data FreeVarEntry rep = FreeVarEntry
{ forall rep. FreeVarEntry rep -> NameInfo rep
freeVarDec :: NameInfo rep,
forall rep. FreeVarEntry rep -> Names
freeVarAliases :: Names,
forall rep. FreeVarEntry rep -> VName -> IndexArray
freeVarIndex :: VName -> IndexArray
}
instance ASTRep rep => Typed (Entry rep) where
typeOf :: Entry rep -> Type
typeOf = forall t. Typed t => t -> Type
typeOf forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Entry rep -> NameInfo rep
entryInfo
entryInfo :: Entry rep -> NameInfo rep
entryInfo :: forall rep. Entry rep -> NameInfo rep
entryInfo Entry rep
e = case forall rep. Entry rep -> EntryType rep
entryType Entry rep
e of
LetBound LetBoundEntry rep
entry -> forall rep. LetDec rep -> NameInfo rep
LetName forall a b. (a -> b) -> a -> b
$ forall rep. LetBoundEntry rep -> LetDec rep
letBoundDec LetBoundEntry rep
entry
LoopVar LoopVarEntry rep
entry -> forall rep. IntType -> NameInfo rep
IndexName forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). LoopVarEntry rep -> IntType
loopVarType LoopVarEntry rep
entry
FParam FParamEntry rep
entry -> forall rep. FParamInfo rep -> NameInfo rep
FParamName forall a b. (a -> b) -> a -> b
$ forall rep. FParamEntry rep -> FParamInfo rep
fparamDec FParamEntry rep
entry
LParam LParamEntry rep
entry -> forall rep. LParamInfo rep -> NameInfo rep
LParamName forall a b. (a -> b) -> a -> b
$ forall rep. LParamEntry rep -> LParamInfo rep
lparamDec LParamEntry rep
entry
FreeVar FreeVarEntry rep
entry -> forall rep. FreeVarEntry rep -> NameInfo rep
freeVarDec FreeVarEntry rep
entry
isLetBound :: Entry rep -> Maybe (LetBoundEntry rep)
isLetBound :: forall rep. Entry rep -> Maybe (LetBoundEntry rep)
isLetBound Entry rep
e = case forall rep. Entry rep -> EntryType rep
entryType Entry rep
e of
LetBound LetBoundEntry rep
entry -> forall a. a -> Maybe a
Just LetBoundEntry rep
entry
EntryType rep
_ -> forall a. Maybe a
Nothing
entryStm :: Entry rep -> Maybe (Stm rep)
entryStm :: forall rep. Entry rep -> Maybe (Stm rep)
entryStm = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall rep. LetBoundEntry rep -> Stm rep
letBoundStm forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Entry rep -> Maybe (LetBoundEntry rep)
isLetBound
entryFParam :: Entry rep -> Maybe (FParamInfo rep)
entryFParam :: forall rep. Entry rep -> Maybe (FParamInfo rep)
entryFParam Entry rep
e = case forall rep. Entry rep -> EntryType rep
entryType Entry rep
e of
FParam FParamEntry rep
e' -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall rep. FParamEntry rep -> FParamInfo rep
fparamDec FParamEntry rep
e'
EntryType rep
_ -> forall a. Maybe a
Nothing
entryLParam :: Entry rep -> Maybe (LParamInfo rep)
entryLParam :: forall rep. Entry rep -> Maybe (LParamInfo rep)
entryLParam Entry rep
e = case forall rep. Entry rep -> EntryType rep
entryType Entry rep
e of
LParam LParamEntry rep
e' -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall rep. LParamEntry rep -> LParamInfo rep
lparamDec LParamEntry rep
e'
EntryType rep
_ -> forall a. Maybe a
Nothing
entryLetBoundDec :: Entry rep -> Maybe (LetDec rep)
entryLetBoundDec :: forall rep. Entry rep -> Maybe (LetDec rep)
entryLetBoundDec = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall rep. LetBoundEntry rep -> LetDec rep
letBoundDec forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Entry rep -> Maybe (LetBoundEntry rep)
isLetBound
entryAliases :: EntryType rep -> Names
entryAliases :: forall rep. EntryType rep -> Names
entryAliases (LetBound LetBoundEntry rep
e) = forall rep. LetBoundEntry rep -> Names
letBoundAliases LetBoundEntry rep
e
entryAliases (FParam FParamEntry rep
e) = forall rep. FParamEntry rep -> Names
fparamAliases FParamEntry rep
e
entryAliases (LParam LParamEntry rep
e) = forall rep. LParamEntry rep -> Names
lparamAliases LParamEntry rep
e
entryAliases (FreeVar FreeVarEntry rep
e) = forall rep. FreeVarEntry rep -> Names
freeVarAliases FreeVarEntry rep
e
entryAliases (LoopVar LoopVarEntry rep
_) = forall a. Monoid a => a
mempty
elem :: VName -> SymbolTable rep -> Bool
elem :: forall rep. VName -> SymbolTable rep -> Bool
elem VName
name = forall a. Maybe a -> Bool
isJust forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. VName -> SymbolTable rep -> Maybe (Entry rep)
lookup VName
name
lookup :: VName -> SymbolTable rep -> Maybe (Entry rep)
lookup :: forall rep. VName -> SymbolTable rep -> Maybe (Entry rep)
lookup VName
name = forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings
lookupStm :: VName -> SymbolTable rep -> Maybe (Stm rep)
lookupStm :: forall rep. VName -> SymbolTable rep -> Maybe (Stm rep)
lookupStm VName
name SymbolTable rep
vtable = forall rep. Entry rep -> Maybe (Stm rep)
entryStm forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall rep. VName -> SymbolTable rep -> Maybe (Entry rep)
lookup VName
name SymbolTable rep
vtable
lookupExp :: VName -> SymbolTable rep -> Maybe (Exp rep, Certs)
lookupExp :: forall rep. VName -> SymbolTable rep -> Maybe (Exp rep, Certs)
lookupExp VName
name SymbolTable rep
vtable = (forall rep. Stm rep -> Exp rep
stmExp forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& forall rep. Stm rep -> Certs
stmCerts) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep. VName -> SymbolTable rep -> Maybe (Stm rep)
lookupStm VName
name SymbolTable rep
vtable
lookupBasicOp :: VName -> SymbolTable rep -> Maybe (BasicOp, Certs)
lookupBasicOp :: forall rep. VName -> SymbolTable rep -> Maybe (BasicOp, Certs)
lookupBasicOp VName
name SymbolTable rep
vtable = case forall rep. VName -> SymbolTable rep -> Maybe (Exp rep, Certs)
lookupExp VName
name SymbolTable rep
vtable of
Just (BasicOp BasicOp
e, Certs
cs) -> forall a. a -> Maybe a
Just (BasicOp
e, Certs
cs)
Maybe (Exp rep, Certs)
_ -> forall a. Maybe a
Nothing
lookupType :: ASTRep rep => VName -> SymbolTable rep -> Maybe Type
lookupType :: forall rep. ASTRep rep => VName -> SymbolTable rep -> Maybe Type
lookupType VName
name SymbolTable rep
vtable = forall t. Typed t => t -> Type
typeOf forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep. VName -> SymbolTable rep -> Maybe (Entry rep)
lookup VName
name SymbolTable rep
vtable
lookupSubExpType :: ASTRep rep => SubExp -> SymbolTable rep -> Maybe Type
lookupSubExpType :: forall rep. ASTRep rep => SubExp -> SymbolTable rep -> Maybe Type
lookupSubExpType (Var VName
v) = forall rep. ASTRep rep => VName -> SymbolTable rep -> Maybe Type
lookupType VName
v
lookupSubExpType (Constant PrimValue
v) = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall shape u. PrimType -> TypeBase shape u
Prim forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
v
lookupSubExp :: VName -> SymbolTable rep -> Maybe (SubExp, Certs)
lookupSubExp :: forall rep. VName -> SymbolTable rep -> Maybe (SubExp, Certs)
lookupSubExp VName
name SymbolTable rep
vtable = do
(Exp rep
e, Certs
cs) <- forall rep. VName -> SymbolTable rep -> Maybe (Exp rep, Certs)
lookupExp VName
name SymbolTable rep
vtable
case Exp rep
e of
BasicOp (SubExp SubExp
se) -> forall a. a -> Maybe a
Just (SubExp
se, Certs
cs)
Exp rep
_ -> forall a. Maybe a
Nothing
lookupAliases :: VName -> SymbolTable rep -> Names
lookupAliases :: forall rep. VName -> SymbolTable rep -> Names
lookupAliases VName
name SymbolTable rep
vtable =
forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty (forall rep. EntryType rep -> Names
entryAliases forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Entry rep -> EntryType rep
entryType) forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable)
lookupLoopVar :: VName -> SymbolTable rep -> Maybe SubExp
lookupLoopVar :: forall rep. VName -> SymbolTable rep -> Maybe SubExp
lookupLoopVar VName
name SymbolTable rep
vtable = do
LoopVar LoopVarEntry rep
e <- forall rep. Entry rep -> EntryType rep
entryType forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable)
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). LoopVarEntry rep -> SubExp
loopVarBound LoopVarEntry rep
e
lookupLoopParam :: VName -> SymbolTable rep -> Maybe (SubExp, SubExp)
lookupLoopParam :: forall rep. VName -> SymbolTable rep -> Maybe (SubExp, SubExp)
lookupLoopParam VName
name SymbolTable rep
vtable = do
FParam FParamEntry rep
e <- forall rep. Entry rep -> EntryType rep
entryType forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable)
forall rep. FParamEntry rep -> Maybe (SubExp, SubExp)
fparamMerge FParamEntry rep
e
aliases :: VName -> VName -> SymbolTable rep -> Bool
aliases :: forall rep. VName -> VName -> SymbolTable rep -> Bool
aliases VName
x VName
y SymbolTable rep
vtable = VName
x forall a. Eq a => a -> a -> Bool
== VName
y Bool -> Bool -> Bool
|| (VName
x VName -> Names -> Bool
`nameIn` forall rep. VName -> SymbolTable rep -> Names
lookupAliases VName
y SymbolTable rep
vtable)
available :: VName -> SymbolTable rep -> Bool
available :: forall rep. VName -> SymbolTable rep -> Bool
available VName
name = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Entry rep -> Bool
entryConsumed) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings
subExpAvailable :: SubExp -> SymbolTable rep -> Bool
subExpAvailable :: forall rep. SubExp -> SymbolTable rep -> Bool
subExpAvailable (Var VName
name) = forall rep. VName -> SymbolTable rep -> Bool
available VName
name
subExpAvailable Constant {} = forall a b. a -> b -> a
const Bool
True
index ::
ASTRep rep =>
VName ->
[SubExp] ->
SymbolTable rep ->
Maybe Indexed
index :: forall rep.
ASTRep rep =>
VName -> [SubExp] -> SymbolTable rep -> Maybe Indexed
index VName
name [SubExp]
is SymbolTable rep
table = do
[TPrimExp Int64 VName]
is' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> Maybe (TPrimExp Int64 VName)
asPrimExp [SubExp]
is
forall rep.
VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
index' VName
name [TPrimExp Int64 VName]
is' SymbolTable rep
table
where
asPrimExp :: SubExp -> Maybe (TPrimExp Int64 VName)
asPrimExp SubExp
i = do
Prim PrimType
t <- forall rep. ASTRep rep => SubExp -> SymbolTable rep -> Maybe Type
lookupSubExpType SubExp
i SymbolTable rep
table
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k) v. PrimExp v -> TPrimExp t v
TPrimExp forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
t SubExp
i
index' ::
VName ->
[TPrimExp Int64 VName] ->
SymbolTable rep ->
Maybe Indexed
index' :: forall rep.
VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
index' VName
name [TPrimExp Int64 VName]
is SymbolTable rep
vtable = do
Entry rep
entry <- forall rep. VName -> SymbolTable rep -> Maybe (Entry rep)
lookup VName
name SymbolTable rep
vtable
case forall rep. Entry rep -> EntryType rep
entryType Entry rep
entry of
LetBound LetBoundEntry rep
entry'
| Just Int
k <-
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex VName
name forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. Pat dec -> [VName]
patNames forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Stm rep -> Pat (LetDec rep)
stmPat forall a b. (a -> b) -> a -> b
$
forall rep. LetBoundEntry rep -> Stm rep
letBoundStm LetBoundEntry rep
entry' ->
forall rep. LetBoundEntry rep -> Int -> IndexArray
letBoundIndex LetBoundEntry rep
entry' Int
k [TPrimExp Int64 VName]
is
FreeVar FreeVarEntry rep
entry' ->
forall rep. FreeVarEntry rep -> VName -> IndexArray
freeVarIndex FreeVarEntry rep
entry' VName
name [TPrimExp Int64 VName]
is
LParam LParamEntry rep
entry' -> forall rep. LParamEntry rep -> IndexArray
lparamIndex LParamEntry rep
entry' [TPrimExp Int64 VName]
is
EntryType rep
_ -> forall a. Maybe a
Nothing
class IndexOp op where
indexOp ::
(ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep ->
Int ->
op ->
[TPrimExp Int64 VName] ->
Maybe Indexed
indexOp SymbolTable rep
_ Int
_ op
_ [TPrimExp Int64 VName]
_ = forall a. Maybe a
Nothing
instance IndexOp (NoOp rep)
indexExp ::
(IndexOp (Op rep), ASTRep rep) =>
SymbolTable rep ->
Exp rep ->
Int ->
IndexArray
indexExp :: forall rep.
(IndexOp (Op rep), ASTRep rep) =>
SymbolTable rep -> Exp rep -> Int -> IndexArray
indexExp SymbolTable rep
vtable (Op Op rep
op) Int
k [TPrimExp Int64 VName]
is =
forall op rep.
(IndexOp op, ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep -> Int -> op -> IndexArray
indexOp SymbolTable rep
vtable Int
k Op rep
op [TPrimExp Int64 VName]
is
indexExp SymbolTable rep
_ (BasicOp (Iota SubExp
_ SubExp
x SubExp
s IntType
to_it)) Int
_ [TPrimExp Int64 VName
i] =
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
Certs -> PrimExp VName -> Indexed
Indexed forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
( forall v. IntType -> PrimExp v -> PrimExp v
sExt IntType
to_it (forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped TPrimExp Int64 VName
i)
PrimExp VName -> PrimExp VName -> PrimExp VName
`mul` PrimType -> SubExp -> PrimExp VName
primExpFromSubExp (IntType -> PrimType
IntType IntType
to_it) SubExp
s
)
PrimExp VName -> PrimExp VName -> PrimExp VName
`add` PrimType -> SubExp -> PrimExp VName
primExpFromSubExp (IntType -> PrimType
IntType IntType
to_it) SubExp
x
where
mul :: PrimExp VName -> PrimExp VName -> PrimExp VName
mul = forall v. BinOp -> PrimExp v -> PrimExp v -> PrimExp v
BinOpExp (IntType -> Overflow -> BinOp
Mul IntType
to_it Overflow
OverflowWrap)
add :: PrimExp VName -> PrimExp VName -> PrimExp VName
add = forall v. BinOp -> PrimExp v -> PrimExp v -> PrimExp v
BinOpExp (IntType -> Overflow -> BinOp
Add IntType
to_it Overflow
OverflowWrap)
indexExp SymbolTable rep
table (BasicOp (Replicate (Shape [SubExp]
ds) SubExp
v)) Int
_ [TPrimExp Int64 VName]
is
| forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
ds forall a. Eq a => a -> a -> Bool
== forall (t :: * -> *) a. Foldable t => t a -> Int
length [TPrimExp Int64 VName]
is,
Just (Prim PrimType
t) <- forall rep. ASTRep rep => SubExp -> SymbolTable rep -> Maybe Type
lookupSubExpType SubExp
v SymbolTable rep
table =
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Certs -> PrimExp VName -> Indexed
Indexed forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
t SubExp
v
indexExp SymbolTable rep
table (BasicOp (Replicate (Shape [SubExp
_]) (Var VName
v))) Int
_ (TPrimExp Int64 VName
_ : [TPrimExp Int64 VName]
is) = do
forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ VName
v forall rep. VName -> SymbolTable rep -> Bool
`available` SymbolTable rep
table
forall rep.
VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
index' VName
v [TPrimExp Int64 VName]
is SymbolTable rep
table
indexExp SymbolTable rep
table (BasicOp (Reshape ReshapeKind
_ ShapeBase SubExp
newshape VName
v)) Int
_ [TPrimExp Int64 VName]
is
| Just [SubExp]
oldshape <- forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
arrayDims forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep. ASTRep rep => VName -> SymbolTable rep -> Maybe Type
lookupType VName
v SymbolTable rep
table =
let is' :: [TPrimExp Int64 VName]
is' =
forall num. IntegralExp num => [num] -> [num] -> [num] -> [num]
reshapeIndex
(forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
oldshape)
(forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 forall a b. (a -> b) -> a -> b
$ forall d. ShapeBase d -> [d]
shapeDims ShapeBase SubExp
newshape)
[TPrimExp Int64 VName]
is
in forall rep.
VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
index' VName
v [TPrimExp Int64 VName]
is' SymbolTable rep
table
indexExp SymbolTable rep
table (BasicOp (Index VName
v Slice SubExp
slice)) Int
_ [TPrimExp Int64 VName]
is = do
forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ VName
v forall rep. VName -> SymbolTable rep -> Bool
`available` SymbolTable rep
table
forall rep.
VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
index' VName
v ([DimIndex SubExp]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
adjust (forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slice) [TPrimExp Int64 VName]
is) SymbolTable rep
table
where
adjust :: [DimIndex SubExp]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
adjust (DimFix SubExp
j : [DimIndex SubExp]
js') [TPrimExp Int64 VName]
is' =
SubExp -> TPrimExp Int64 VName
pe64 SubExp
j forall a. a -> [a] -> [a]
: [DimIndex SubExp]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
adjust [DimIndex SubExp]
js' [TPrimExp Int64 VName]
is'
adjust (DimSlice SubExp
j SubExp
_ SubExp
s : [DimIndex SubExp]
js') (TPrimExp Int64 VName
i : [TPrimExp Int64 VName]
is') =
let i_t_s :: TPrimExp Int64 VName
i_t_s = TPrimExp Int64 VName
i forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
s
j_p_i_t_s :: TPrimExp Int64 VName
j_p_i_t_s = SubExp -> TPrimExp Int64 VName
pe64 SubExp
j forall a. Num a => a -> a -> a
+ TPrimExp Int64 VName
i_t_s
in TPrimExp Int64 VName
j_p_i_t_s forall a. a -> [a] -> [a]
: [DimIndex SubExp]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
adjust [DimIndex SubExp]
js' [TPrimExp Int64 VName]
is'
adjust [DimIndex SubExp]
_ [TPrimExp Int64 VName]
_ = []
indexExp SymbolTable rep
_ Exp rep
_ Int
_ [TPrimExp Int64 VName]
_ = forall a. Maybe a
Nothing
defBndEntry ::
(ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep ->
PatElem (LetDec rep) ->
Names ->
Stm rep ->
LetBoundEntry rep
defBndEntry :: forall rep.
(ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep
-> PatElem (LetDec rep) -> Names -> Stm rep -> LetBoundEntry rep
defBndEntry SymbolTable rep
vtable PatElem (LetDec rep)
patElem Names
als Stm rep
stm =
LetBoundEntry
{ letBoundDec :: LetDec rep
letBoundDec = forall dec. PatElem dec -> dec
patElemDec PatElem (LetDec rep)
patElem,
letBoundAliases :: Names
letBoundAliases = Names
als,
letBoundStm :: Stm rep
letBoundStm = Stm rep
stm,
letBoundIndex :: Int -> IndexArray
letBoundIndex = \Int
k ->
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Certs -> Indexed -> Indexed
indexedAddCerts (forall dec. StmAux dec -> Certs
stmAuxCerts forall a b. (a -> b) -> a -> b
$ forall rep. Stm rep -> StmAux (ExpDec rep)
stmAux Stm rep
stm))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep.
(IndexOp (Op rep), ASTRep rep) =>
SymbolTable rep -> Exp rep -> Int -> IndexArray
indexExp SymbolTable rep
vtable (forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm) Int
k
}
bindingEntries ::
(Aliases.Aliased rep, IndexOp (Op rep)) =>
Stm rep ->
SymbolTable rep ->
[LetBoundEntry rep]
bindingEntries :: forall rep.
(Aliased rep, IndexOp (Op rep)) =>
Stm rep -> SymbolTable rep -> [LetBoundEntry rep]
bindingEntries stm :: Stm rep
stm@(Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
_ Exp rep
_) SymbolTable rep
vtable = do
PatElem (LetDec rep)
pat_elem <- forall dec. Pat dec -> [PatElem dec]
patElems Pat (LetDec rep)
pat
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall rep.
(ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep
-> PatElem (LetDec rep) -> Names -> Stm rep -> LetBoundEntry rep
defBndEntry SymbolTable rep
vtable PatElem (LetDec rep)
pat_elem (forall a. AliasesOf a => a -> Names
Aliases.aliasesOf PatElem (LetDec rep)
pat_elem) Stm rep
stm
adjustSeveral :: Ord k => (v -> v) -> [k] -> M.Map k v -> M.Map k v
adjustSeveral :: forall k v. Ord k => (v -> v) -> [k] -> Map k v -> Map k v
adjustSeveral v -> v
f = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust v -> v
f
insertEntry ::
ASTRep rep =>
VName ->
EntryType rep ->
SymbolTable rep ->
SymbolTable rep
insertEntry :: forall rep.
ASTRep rep =>
VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
insertEntry VName
name EntryType rep
entry SymbolTable rep
vtable =
let entry' :: Entry rep
entry' =
Entry
{ entryConsumed :: Bool
entryConsumed = Bool
False,
entryDepth :: Int
entryDepth = forall rep. SymbolTable rep -> Int
loopDepth SymbolTable rep
vtable,
entryIsSize :: Bool
entryIsSize = Bool
False,
entryAccInput :: Maybe (WithAccInput rep)
entryAccInput = forall a. Maybe a
Nothing,
entryType :: EntryType rep
entryType = EntryType rep
entry
}
dims :: [VName]
dims = forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe SubExp -> Maybe VName
subExpVar forall a b. (a -> b) -> a -> b
$ forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
arrayDims forall a b. (a -> b) -> a -> b
$ forall t. Typed t => t -> Type
typeOf Entry rep
entry'
isSize :: Entry rep -> Entry rep
isSize Entry rep
e = Entry rep
e {entryIsSize :: Bool
entryIsSize = Bool
True}
in SymbolTable rep
vtable
{ bindings :: Map VName (Entry rep)
bindings =
forall k v. Ord k => (v -> v) -> [k] -> Map k v -> Map k v
adjustSeveral forall {rep}. Entry rep -> Entry rep
isSize [VName]
dims forall a b. (a -> b) -> a -> b
$
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
name Entry rep
entry' forall a b. (a -> b) -> a -> b
$
forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable
}
insertEntries ::
ASTRep rep =>
[(VName, EntryType rep)] ->
SymbolTable rep ->
SymbolTable rep
insertEntries :: forall rep.
ASTRep rep =>
[(VName, EntryType rep)] -> SymbolTable rep -> SymbolTable rep
insertEntries [(VName, EntryType rep)]
entries SymbolTable rep
vtable =
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {rep}.
ASTRep rep =>
SymbolTable rep -> (VName, EntryType rep) -> SymbolTable rep
add SymbolTable rep
vtable [(VName, EntryType rep)]
entries
where
add :: SymbolTable rep -> (VName, EntryType rep) -> SymbolTable rep
add SymbolTable rep
vtable' (VName
name, EntryType rep
entry) = forall rep.
ASTRep rep =>
VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
insertEntry VName
name EntryType rep
entry SymbolTable rep
vtable'
insertStm ::
(IndexOp (Op rep), Aliases.Aliased rep) =>
Stm rep ->
SymbolTable rep ->
SymbolTable rep
insertStm :: forall rep.
(IndexOp (Op rep), Aliased rep) =>
Stm rep -> SymbolTable rep -> SymbolTable rep
insertStm Stm rep
stm SymbolTable rep
vtable =
forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip forall rep. VName -> SymbolTable rep -> SymbolTable rep
consume) (Names -> [VName]
namesToList Names
stm_consumed) forall a b. (a -> b) -> a -> b
$
forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {dec} {rep}.
AliasesOf dec =>
SymbolTable rep -> PatElem dec -> SymbolTable rep
addRevAliases) (forall dec. Pat dec -> [PatElem dec]
patElems forall a b. (a -> b) -> a -> b
$ forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm rep
stm) forall a b. (a -> b) -> a -> b
$
forall rep.
ASTRep rep =>
[(VName, EntryType rep)] -> SymbolTable rep -> SymbolTable rep
insertEntries (forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
names forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall rep. LetBoundEntry rep -> EntryType rep
LetBound forall a b. (a -> b) -> a -> b
$ forall rep.
(Aliased rep, IndexOp (Op rep)) =>
Stm rep -> SymbolTable rep -> [LetBoundEntry rep]
bindingEntries Stm rep
stm SymbolTable rep
vtable) SymbolTable rep
vtable
where
names :: [VName]
names = forall dec. Pat dec -> [VName]
patNames forall a b. (a -> b) -> a -> b
$ forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm rep
stm
stm_consumed :: Names
stm_consumed = forall rep. Names -> SymbolTable rep -> Names
expandAliases (forall rep. Aliased rep => Stm rep -> Names
Aliases.consumedInStm Stm rep
stm) SymbolTable rep
vtable
addRevAliases :: SymbolTable rep -> PatElem dec -> SymbolTable rep
addRevAliases SymbolTable rep
vtable' PatElem dec
pe =
SymbolTable rep
vtable' {bindings :: Map VName (Entry rep)
bindings = forall k v. Ord k => (v -> v) -> [k] -> Map k v -> Map k v
adjustSeveral Entry rep -> Entry rep
update [VName]
inedges forall a b. (a -> b) -> a -> b
$ forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable'}
where
inedges :: [VName]
inedges = Names -> [VName]
namesToList forall a b. (a -> b) -> a -> b
$ forall rep. Names -> SymbolTable rep -> Names
expandAliases (forall a. AliasesOf a => a -> Names
Aliases.aliasesOf PatElem dec
pe) SymbolTable rep
vtable'
update :: Entry rep -> Entry rep
update Entry rep
e = Entry rep
e {entryType :: EntryType rep
entryType = EntryType rep -> EntryType rep
update' forall a b. (a -> b) -> a -> b
$ forall rep. Entry rep -> EntryType rep
entryType Entry rep
e}
update' :: EntryType rep -> EntryType rep
update' (LetBound LetBoundEntry rep
entry) =
forall rep. LetBoundEntry rep -> EntryType rep
LetBound
LetBoundEntry rep
entry
{ letBoundAliases :: Names
letBoundAliases = VName -> Names
oneName (forall dec. PatElem dec -> VName
patElemName PatElem dec
pe) forall a. Semigroup a => a -> a -> a
<> forall rep. LetBoundEntry rep -> Names
letBoundAliases LetBoundEntry rep
entry
}
update' (FParam FParamEntry rep
entry) =
forall rep. FParamEntry rep -> EntryType rep
FParam
FParamEntry rep
entry
{ fparamAliases :: Names
fparamAliases = VName -> Names
oneName (forall dec. PatElem dec -> VName
patElemName PatElem dec
pe) forall a. Semigroup a => a -> a -> a
<> forall rep. FParamEntry rep -> Names
fparamAliases FParamEntry rep
entry
}
update' (LParam LParamEntry rep
entry) =
forall rep. LParamEntry rep -> EntryType rep
LParam
LParamEntry rep
entry
{ lparamAliases :: Names
lparamAliases = VName -> Names
oneName (forall dec. PatElem dec -> VName
patElemName PatElem dec
pe) forall a. Semigroup a => a -> a -> a
<> forall rep. LParamEntry rep -> Names
lparamAliases LParamEntry rep
entry
}
update' (FreeVar FreeVarEntry rep
entry) =
forall rep. FreeVarEntry rep -> EntryType rep
FreeVar
FreeVarEntry rep
entry
{ freeVarAliases :: Names
freeVarAliases = VName -> Names
oneName (forall dec. PatElem dec -> VName
patElemName PatElem dec
pe) forall a. Semigroup a => a -> a -> a
<> forall rep. FreeVarEntry rep -> Names
freeVarAliases FreeVarEntry rep
entry
}
update' EntryType rep
e = EntryType rep
e
insertStms ::
(IndexOp (Op rep), Aliases.Aliased rep) =>
Stms rep ->
SymbolTable rep ->
SymbolTable rep
insertStms :: forall rep.
(IndexOp (Op rep), Aliased rep) =>
Stms rep -> SymbolTable rep -> SymbolTable rep
insertStms Stms rep
stms SymbolTable rep
vtable = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall rep.
(IndexOp (Op rep), Aliased rep) =>
Stm rep -> SymbolTable rep -> SymbolTable rep
insertStm) SymbolTable rep
vtable forall a b. (a -> b) -> a -> b
$ forall rep. Stms rep -> [Stm rep]
stmsToList Stms rep
stms
expandAliases :: Names -> SymbolTable rep -> Names
expandAliases :: forall rep. Names -> SymbolTable rep -> Names
expandAliases Names
names SymbolTable rep
vtable = Names
names forall a. Semigroup a => a -> a -> a
<> Names
aliasesOfAliases
where
aliasesOfAliases :: Names
aliasesOfAliases =
forall a. Monoid a => [a] -> a
mconcat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall rep. VName -> SymbolTable rep -> Names
`lookupAliases` SymbolTable rep
vtable) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Names -> [VName]
namesToList forall a b. (a -> b) -> a -> b
$ Names
names
insertFParam ::
ASTRep rep =>
AST.FParam rep ->
SymbolTable rep ->
SymbolTable rep
insertFParam :: forall rep.
ASTRep rep =>
FParam rep -> SymbolTable rep -> SymbolTable rep
insertFParam FParam rep
fparam = forall rep.
ASTRep rep =>
VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
insertEntry VName
name EntryType rep
entry
where
name :: VName
name = forall dec. Param dec -> VName
AST.paramName FParam rep
fparam
entry :: EntryType rep
entry =
forall rep. FParamEntry rep -> EntryType rep
FParam
FParamEntry
{ fparamDec :: FParamInfo rep
fparamDec = forall dec. Param dec -> dec
AST.paramDec FParam rep
fparam,
fparamAliases :: Names
fparamAliases = forall a. Monoid a => a
mempty,
fparamMerge :: Maybe (SubExp, SubExp)
fparamMerge = forall a. Maybe a
Nothing
}
insertFParams ::
ASTRep rep =>
[AST.FParam rep] ->
SymbolTable rep ->
SymbolTable rep
insertFParams :: forall rep.
ASTRep rep =>
[FParam rep] -> SymbolTable rep -> SymbolTable rep
insertFParams [FParam rep]
fparams SymbolTable rep
symtable = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall rep.
ASTRep rep =>
FParam rep -> SymbolTable rep -> SymbolTable rep
insertFParam) SymbolTable rep
symtable [FParam rep]
fparams
insertLParam :: ASTRep rep => LParam rep -> SymbolTable rep -> SymbolTable rep
insertLParam :: forall rep.
ASTRep rep =>
LParam rep -> SymbolTable rep -> SymbolTable rep
insertLParam LParam rep
param = forall rep.
ASTRep rep =>
VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
insertEntry VName
name EntryType rep
bind
where
bind :: EntryType rep
bind =
forall rep. LParamEntry rep -> EntryType rep
LParam
LParamEntry
{ lparamDec :: LParamInfo rep
lparamDec = forall dec. Param dec -> dec
AST.paramDec LParam rep
param,
lparamAliases :: Names
lparamAliases = forall a. Monoid a => a
mempty,
lparamIndex :: IndexArray
lparamIndex = forall a b. a -> b -> a
const forall a. Maybe a
Nothing
}
name :: VName
name = forall dec. Param dec -> VName
AST.paramName LParam rep
param
insertLoopMerge ::
ASTRep rep =>
[(AST.FParam rep, SubExp, SubExpRes)] ->
SymbolTable rep ->
SymbolTable rep
insertLoopMerge :: forall rep.
ASTRep rep =>
[(FParam rep, SubExp, SubExpRes)]
-> SymbolTable rep -> SymbolTable rep
insertLoopMerge = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip forall {rep}.
ASTRep rep =>
(Param (FParamInfo rep), SubExp, SubExpRes)
-> SymbolTable rep -> SymbolTable rep
bind
where
bind :: (Param (FParamInfo rep), SubExp, SubExpRes)
-> SymbolTable rep -> SymbolTable rep
bind (Param (FParamInfo rep)
p, SubExp
initial, SubExpRes Certs
_ SubExp
res) =
forall rep.
ASTRep rep =>
VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
insertEntry (forall dec. Param dec -> VName
paramName Param (FParamInfo rep)
p) forall a b. (a -> b) -> a -> b
$
forall rep. FParamEntry rep -> EntryType rep
FParam
FParamEntry
{ fparamDec :: FParamInfo rep
fparamDec = forall dec. Param dec -> dec
AST.paramDec Param (FParamInfo rep)
p,
fparamAliases :: Names
fparamAliases = forall a. Monoid a => a
mempty,
fparamMerge :: Maybe (SubExp, SubExp)
fparamMerge = forall a. a -> Maybe a
Just (SubExp
initial, SubExp
res)
}
insertLoopVar :: ASTRep rep => VName -> IntType -> SubExp -> SymbolTable rep -> SymbolTable rep
insertLoopVar :: forall rep.
ASTRep rep =>
VName -> IntType -> SubExp -> SymbolTable rep -> SymbolTable rep
insertLoopVar VName
name IntType
it SubExp
bound = forall rep.
ASTRep rep =>
VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
insertEntry VName
name EntryType rep
bind
where
bind :: EntryType rep
bind =
forall rep. LoopVarEntry rep -> EntryType rep
LoopVar
LoopVarEntry
{ loopVarType :: IntType
loopVarType = IntType
it,
loopVarBound :: SubExp
loopVarBound = SubExp
bound
}
insertFreeVar :: ASTRep rep => VName -> NameInfo rep -> SymbolTable rep -> SymbolTable rep
insertFreeVar :: forall rep.
ASTRep rep =>
VName -> NameInfo rep -> SymbolTable rep -> SymbolTable rep
insertFreeVar VName
name NameInfo rep
dec = forall rep.
ASTRep rep =>
VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
insertEntry VName
name EntryType rep
entry
where
entry :: EntryType rep
entry =
forall rep. FreeVarEntry rep -> EntryType rep
FreeVar
FreeVarEntry
{ freeVarDec :: NameInfo rep
freeVarDec = NameInfo rep
dec,
freeVarIndex :: VName -> IndexArray
freeVarIndex = \VName
_ [TPrimExp Int64 VName]
_ -> forall a. Maybe a
Nothing,
freeVarAliases :: Names
freeVarAliases = forall a. Monoid a => a
mempty
}
consume :: VName -> SymbolTable rep -> SymbolTable rep
consume :: forall rep. VName -> SymbolTable rep -> SymbolTable rep
consume VName
consumee SymbolTable rep
vtable =
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {rep}. SymbolTable rep -> VName -> SymbolTable rep
consume' SymbolTable rep
vtable forall a b. (a -> b) -> a -> b
$
Names -> [VName]
namesToList forall a b. (a -> b) -> a -> b
$
forall rep. Names -> SymbolTable rep -> Names
expandAliases (VName -> Names
oneName VName
consumee) SymbolTable rep
vtable
where
consume' :: SymbolTable rep -> VName -> SymbolTable rep
consume' SymbolTable rep
vtable' VName
v =
SymbolTable rep
vtable' {bindings :: Map VName (Entry rep)
bindings = forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust forall {rep}. Entry rep -> Entry rep
consume'' VName
v forall a b. (a -> b) -> a -> b
$ forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable'}
consume'' :: Entry rep -> Entry rep
consume'' Entry rep
e = Entry rep
e {entryConsumed :: Bool
entryConsumed = Bool
True}
hideIf :: (Entry rep -> Bool) -> SymbolTable rep -> SymbolTable rep
hideIf :: forall rep.
(Entry rep -> Bool) -> SymbolTable rep -> SymbolTable rep
hideIf Entry rep -> Bool
hide SymbolTable rep
vtable = SymbolTable rep
vtable {bindings :: Map VName (Entry rep)
bindings = forall a b k. (a -> b) -> Map k a -> Map k b
M.map Entry rep -> Entry rep
maybeHide forall a b. (a -> b) -> a -> b
$ forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable}
where
maybeHide :: Entry rep -> Entry rep
maybeHide Entry rep
entry
| Entry rep -> Bool
hide Entry rep
entry =
Entry rep
entry
{ entryType :: EntryType rep
entryType =
forall rep. FreeVarEntry rep -> EntryType rep
FreeVar
FreeVarEntry
{ freeVarDec :: NameInfo rep
freeVarDec = forall rep. Entry rep -> NameInfo rep
entryInfo Entry rep
entry,
freeVarIndex :: VName -> IndexArray
freeVarIndex = \VName
_ [TPrimExp Int64 VName]
_ -> forall a. Maybe a
Nothing,
freeVarAliases :: Names
freeVarAliases = forall rep. EntryType rep -> Names
entryAliases forall a b. (a -> b) -> a -> b
$ forall rep. Entry rep -> EntryType rep
entryType Entry rep
entry
}
}
| Bool
otherwise = Entry rep
entry
hideCertified :: Names -> SymbolTable rep -> SymbolTable rep
hideCertified :: forall rep. Names -> SymbolTable rep -> SymbolTable rep
hideCertified Names
to_hide = forall rep.
(Entry rep -> Bool) -> SymbolTable rep -> SymbolTable rep
hideIf forall a b. (a -> b) -> a -> b
$ forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False Stm rep -> Bool
hide forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Entry rep -> Maybe (Stm rep)
entryStm
where
hide :: Stm rep -> Bool
hide = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Names -> Bool
`nameIn` Names
to_hide) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Certs -> [VName]
unCerts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Stm rep -> Certs
stmCerts
noteAccTokens ::
[(VName, WithAccInput rep)] ->
SymbolTable rep ->
SymbolTable rep
noteAccTokens :: forall rep.
[(VName, WithAccInput rep)] -> SymbolTable rep -> SymbolTable rep
noteAccTokens = forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {rep}.
SymbolTable rep -> (VName, WithAccInput rep) -> SymbolTable rep
f)
where
f :: SymbolTable rep -> (VName, WithAccInput rep) -> SymbolTable rep
f SymbolTable rep
vtable (VName
v, WithAccInput rep
accum) =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v forall a b. (a -> b) -> a -> b
$ forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable of
Maybe (Entry rep)
Nothing -> SymbolTable rep
vtable
Just Entry rep
e ->
SymbolTable rep
vtable
{ bindings :: Map VName (Entry rep)
bindings =
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v (Entry rep
e {entryAccInput :: Maybe (WithAccInput rep)
entryAccInput = forall a. a -> Maybe a
Just WithAccInput rep
accum}) forall a b. (a -> b) -> a -> b
$ forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable
}