{-# LANGUAGE Strict #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
module Futhark.Optimise.Simplify.Engine
(
SimpleM,
runSimpleM,
SimpleOps (..),
SimplifyOp,
bindableSimpleOps,
Env (envHoistBlockers, envRules),
emptyEnv,
HoistBlockers (..),
neverBlocks,
noExtraHoistBlockers,
neverHoist,
BlockPred,
orIf,
hasFree,
isConsumed,
isConsuming,
isFalse,
isOp,
isNotSafe,
isDeviceMigrated,
asksEngineEnv,
askVtable,
localVtable,
SimplifiableRep,
Simplifiable (..),
simplifyFun,
simplifyStms,
simplifyStmsWithUsage,
simplifyLambda,
simplifyLambdaNoHoisting,
bindLParams,
simplifyBody,
ST.SymbolTable,
hoistStms,
blockIf,
blockMigrated,
enterLoop,
constructBody,
module Futhark.Optimise.Simplify.Rep,
)
where
import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.Either
import Data.List (find, foldl', inits, mapAccumL)
import Data.Map qualified as M
import Data.Maybe
import Futhark.Analysis.SymbolTable qualified as ST
import Futhark.Analysis.UsageTable qualified as UT
import Futhark.Construct
import Futhark.IR
import Futhark.IR.Prop.Aliases
import Futhark.Optimise.Simplify.Rep
import Futhark.Optimise.Simplify.Rule
import Futhark.Util (nubOrd)
data HoistBlockers rep = HoistBlockers
{
forall {k} (rep :: k). HoistBlockers rep -> BlockPred (Wise rep)
blockHoistPar :: BlockPred (Wise rep),
forall {k} (rep :: k). HoistBlockers rep -> BlockPred (Wise rep)
blockHoistSeq :: BlockPred (Wise rep),
forall {k} (rep :: k). HoistBlockers rep -> BlockPred (Wise rep)
blockHoistBranch :: BlockPred (Wise rep),
forall {k} (rep :: k). HoistBlockers rep -> Stm (Wise rep) -> Bool
isAllocation :: Stm (Wise rep) -> Bool
}
noExtraHoistBlockers :: HoistBlockers rep
=
forall {k} (rep :: k).
BlockPred (Wise rep)
-> BlockPred (Wise rep)
-> BlockPred (Wise rep)
-> (Stm (Wise rep) -> Bool)
-> HoistBlockers rep
HoistBlockers forall {k} (rep :: k). BlockPred rep
neverBlocks forall {k} (rep :: k). BlockPred rep
neverBlocks forall {k} (rep :: k). BlockPred rep
neverBlocks (forall a b. a -> b -> a
const Bool
False)
neverHoist :: HoistBlockers rep
neverHoist :: forall {k} (rep :: k). HoistBlockers rep
neverHoist =
forall {k} (rep :: k).
BlockPred (Wise rep)
-> BlockPred (Wise rep)
-> BlockPred (Wise rep)
-> (Stm (Wise rep) -> Bool)
-> HoistBlockers rep
HoistBlockers forall {k} (rep :: k). BlockPred rep
alwaysBlocks forall {k} (rep :: k). BlockPred rep
alwaysBlocks forall {k} (rep :: k). BlockPred rep
alwaysBlocks (forall a b. a -> b -> a
const Bool
False)
data Env rep = Env
{ forall {k} (rep :: k). Env rep -> RuleBook (Wise rep)
envRules :: RuleBook (Wise rep),
forall {k} (rep :: k). Env rep -> HoistBlockers rep
envHoistBlockers :: HoistBlockers rep,
forall {k} (rep :: k). Env rep -> SymbolTable (Wise rep)
envVtable :: ST.SymbolTable (Wise rep)
}
emptyEnv :: RuleBook (Wise rep) -> HoistBlockers rep -> Env rep
emptyEnv :: forall {k} (rep :: k).
RuleBook (Wise rep) -> HoistBlockers rep -> Env rep
emptyEnv RuleBook (Wise rep)
rules HoistBlockers rep
blockers =
Env
{ envRules :: RuleBook (Wise rep)
envRules = RuleBook (Wise rep)
rules,
envHoistBlockers :: HoistBlockers rep
envHoistBlockers = HoistBlockers rep
blockers,
envVtable :: SymbolTable (Wise rep)
envVtable = forall a. Monoid a => a
mempty
}
type Protect m = SubExp -> Pat (LetDec (Rep m)) -> Op (Rep m) -> Maybe (m ())
type SimplifyOp rep op = op -> SimpleM rep (op, Stms (Wise rep))
data SimpleOps rep = SimpleOps
{ forall {k} (rep :: k).
SimpleOps rep
-> SymbolTable (Wise rep)
-> Pat (LetDec (Wise rep))
-> Exp (Wise rep)
-> SimpleM rep (ExpDec (Wise rep))
mkExpDecS ::
ST.SymbolTable (Wise rep) ->
Pat (LetDec (Wise rep)) ->
Exp (Wise rep) ->
SimpleM rep (ExpDec (Wise rep)),
forall {k} (rep :: k).
SimpleOps rep
-> SymbolTable (Wise rep)
-> Stms (Wise rep)
-> Result
-> SimpleM rep (Body (Wise rep))
mkBodyS ::
ST.SymbolTable (Wise rep) ->
Stms (Wise rep) ->
Result ->
SimpleM rep (Body (Wise rep)),
forall {k} (rep :: k).
SimpleOps rep -> Protect (Builder (Wise rep))
protectHoistedOpS :: Protect (Builder (Wise rep)),
forall {k} (rep :: k). SimpleOps rep -> Op (Wise rep) -> UsageTable
opUsageS :: Op (Wise rep) -> UT.UsageTable,
forall {k} (rep :: k).
SimpleOps rep
-> Pat (LetDec rep)
-> Exp (Wise rep)
-> SimpleM rep (Pat (LetDec rep))
simplifyPatFromExpS ::
Pat (LetDec rep) ->
Exp (Wise rep) ->
SimpleM rep (Pat (LetDec rep)),
forall {k} (rep :: k).
SimpleOps rep -> SimplifyOp rep (Op (Wise rep))
simplifyOpS :: SimplifyOp rep (Op (Wise rep))
}
bindableSimpleOps ::
(SimplifiableRep rep, Buildable rep) =>
SimplifyOp rep (Op (Wise rep)) ->
SimpleOps rep
bindableSimpleOps :: forall {k} (rep :: k).
(SimplifiableRep rep, Buildable rep) =>
SimplifyOp rep (Op (Wise rep)) -> SimpleOps rep
bindableSimpleOps =
forall {k} (rep :: k).
(SymbolTable (Wise rep)
-> Pat (LetDec (Wise rep))
-> Exp (Wise rep)
-> SimpleM rep (ExpDec (Wise rep)))
-> (SymbolTable (Wise rep)
-> Stms (Wise rep) -> Result -> SimpleM rep (Body (Wise rep)))
-> Protect (Builder (Wise rep))
-> (Op (Wise rep) -> UsageTable)
-> (Pat (LetDec rep)
-> Exp (Wise rep) -> SimpleM rep (Pat (LetDec rep)))
-> SimplifyOp rep (Op (Wise rep))
-> SimpleOps rep
SimpleOps forall {k} {f :: * -> *} {rep :: k} {p}.
(Applicative f, Buildable rep) =>
p -> Pat (LetDec rep) -> Exp rep -> f (ExpDec rep)
mkExpDecS' forall {k} {f :: * -> *} {rep :: k} {p}.
(Applicative f, Buildable rep) =>
p -> Stms rep -> Result -> f (Body rep)
mkBodyS' forall {p} {p} {p} {a}. p -> p -> p -> Maybe a
protectHoistedOpS' (forall a b. a -> b -> a
const forall a. Monoid a => a
mempty) forall {k} {t :: * -> *} {rep :: k} {b} {p}.
(Traversable t, ASTRep rep, Simplifiable b,
Simplifiable (LetDec rep), Simplifiable (FParamInfo rep),
Simplifiable (LParamInfo rep), Simplifiable (RetType rep),
Simplifiable (BranchType rep), TraverseOpStms (Wise rep),
CanBeWise (Op rep), IndexOp (OpWithWisdom (Op rep)),
BuilderOps (Wise rep)) =>
t b -> p -> SimpleM rep (t b)
simplifyPatFromExp
where
mkExpDecS' :: p -> Pat (LetDec rep) -> Exp rep -> f (ExpDec rep)
mkExpDecS' p
_ Pat (LetDec rep)
pat Exp rep
e = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
Buildable rep =>
Pat (LetDec rep) -> Exp rep -> ExpDec rep
mkExpDec Pat (LetDec rep)
pat Exp rep
e
mkBodyS' :: p -> Stms rep -> Result -> f (Body rep)
mkBodyS' p
_ Stms rep
stms Result
res = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
Buildable rep =>
Stms rep -> Result -> Body rep
mkBody Stms rep
stms Result
res
protectHoistedOpS' :: p -> p -> p -> Maybe a
protectHoistedOpS' p
_ p
_ p
_ = forall a. Maybe a
Nothing
simplifyPatFromExp :: t b -> p -> SimpleM rep (t b)
simplifyPatFromExp t b
pat p
_ = forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify t b
pat
newtype SimpleM rep a
= SimpleM
( ReaderT
(SimpleOps rep, Env rep)
(State (VNameSource, Bool, Certs))
a
)
deriving
( forall a. a -> SimpleM rep a
forall {k} {rep :: k}. Functor (SimpleM rep)
forall k (rep :: k) a. a -> SimpleM rep a
forall k (rep :: k) a b.
SimpleM rep a -> SimpleM rep b -> SimpleM rep a
forall k (rep :: k) a b.
SimpleM rep a -> SimpleM rep b -> SimpleM rep b
forall k (rep :: k) a b.
SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
forall k (rep :: k) a b c.
(a -> b -> c) -> SimpleM rep a -> SimpleM rep b -> SimpleM rep c
forall a b. SimpleM rep a -> SimpleM rep b -> SimpleM rep a
forall a b. SimpleM rep a -> SimpleM rep b -> SimpleM rep b
forall a b. SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
forall a b c.
(a -> b -> c) -> SimpleM rep a -> SimpleM rep b -> SimpleM 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. SimpleM rep a -> SimpleM rep b -> SimpleM rep a
$c<* :: forall k (rep :: k) a b.
SimpleM rep a -> SimpleM rep b -> SimpleM rep a
*> :: forall a b. SimpleM rep a -> SimpleM rep b -> SimpleM rep b
$c*> :: forall k (rep :: k) a b.
SimpleM rep a -> SimpleM rep b -> SimpleM rep b
liftA2 :: forall a b c.
(a -> b -> c) -> SimpleM rep a -> SimpleM rep b -> SimpleM rep c
$cliftA2 :: forall k (rep :: k) a b c.
(a -> b -> c) -> SimpleM rep a -> SimpleM rep b -> SimpleM rep c
<*> :: forall a b. SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
$c<*> :: forall k (rep :: k) a b.
SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
pure :: forall a. a -> SimpleM rep a
$cpure :: forall k (rep :: k) a. a -> SimpleM rep a
Applicative,
forall k (rep :: k) a b. a -> SimpleM rep b -> SimpleM rep a
forall k (rep :: k) a b. (a -> b) -> SimpleM rep a -> SimpleM rep b
forall a b. a -> SimpleM rep b -> SimpleM rep a
forall a b. (a -> b) -> SimpleM rep a -> SimpleM 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 -> SimpleM rep b -> SimpleM rep a
$c<$ :: forall k (rep :: k) a b. a -> SimpleM rep b -> SimpleM rep a
fmap :: forall a b. (a -> b) -> SimpleM rep a -> SimpleM rep b
$cfmap :: forall k (rep :: k) a b. (a -> b) -> SimpleM rep a -> SimpleM rep b
Functor,
forall a. a -> SimpleM rep a
forall k (rep :: k). Applicative (SimpleM rep)
forall k (rep :: k) a. a -> SimpleM rep a
forall k (rep :: k) a b.
SimpleM rep a -> SimpleM rep b -> SimpleM rep b
forall k (rep :: k) a b.
SimpleM rep a -> (a -> SimpleM rep b) -> SimpleM rep b
forall a b. SimpleM rep a -> SimpleM rep b -> SimpleM rep b
forall a b. SimpleM rep a -> (a -> SimpleM rep b) -> SimpleM 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 -> SimpleM rep a
$creturn :: forall k (rep :: k) a. a -> SimpleM rep a
>> :: forall a b. SimpleM rep a -> SimpleM rep b -> SimpleM rep b
$c>> :: forall k (rep :: k) a b.
SimpleM rep a -> SimpleM rep b -> SimpleM rep b
>>= :: forall a b. SimpleM rep a -> (a -> SimpleM rep b) -> SimpleM rep b
$c>>= :: forall k (rep :: k) a b.
SimpleM rep a -> (a -> SimpleM rep b) -> SimpleM rep b
Monad,
MonadReader (SimpleOps rep, Env rep),
MonadState (VNameSource, Bool, Certs)
)
instance MonadFreshNames (SimpleM rep) where
putNameSource :: VNameSource -> SimpleM rep ()
putNameSource VNameSource
src = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \(VNameSource
_, Bool
b, Certs
c) -> (VNameSource
src, Bool
b, Certs
c)
getNameSource :: SimpleM rep VNameSource
getNameSource = forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall a b. (a -> b) -> a -> b
$ \(VNameSource
a, Bool
_, Certs
_) -> VNameSource
a
instance SimplifiableRep rep => HasScope (Wise rep) (SimpleM rep) where
askScope :: SimpleM rep (Scope (Wise rep))
askScope = forall {k} (rep :: k). SymbolTable rep -> Scope rep
ST.toScope forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (rep :: k). SimpleM rep (SymbolTable (Wise rep))
askVtable
lookupType :: VName -> SimpleM rep Type
lookupType VName
name = do
SymbolTable (Wise rep)
vtable <- forall {k} (rep :: k). SimpleM rep (SymbolTable (Wise rep))
askVtable
case forall {k} (rep :: k).
ASTRep rep =>
VName -> SymbolTable rep -> Maybe Type
ST.lookupType VName
name SymbolTable (Wise rep)
vtable of
Just Type
t -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
t
Maybe Type
Nothing ->
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
"SimpleM.lookupType: cannot find variable "
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString VName
name
forall a. [a] -> [a] -> [a]
++ [Char]
" in symbol table."
instance
SimplifiableRep rep =>
LocalScope (Wise rep) (SimpleM rep)
where
localScope :: forall a. Scope (Wise rep) -> SimpleM rep a -> SimpleM rep a
localScope Scope (Wise rep)
types = forall {k} (rep :: k) a.
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep a -> SimpleM rep a
localVtable (forall a. Semigroup a => a -> a -> a
<> forall {k} (rep :: k). ASTRep rep => Scope rep -> SymbolTable rep
ST.fromScope Scope (Wise rep)
types)
runSimpleM ::
SimpleM rep a ->
SimpleOps rep ->
Env rep ->
VNameSource ->
((a, Bool), VNameSource)
runSimpleM :: forall {k} (rep :: k) a.
SimpleM rep a
-> SimpleOps rep
-> Env rep
-> VNameSource
-> ((a, Bool), VNameSource)
runSimpleM (SimpleM ReaderT
(SimpleOps rep, Env rep) (State (VNameSource, Bool, Certs)) a
m) SimpleOps rep
simpl Env rep
env VNameSource
src =
let (a
x, (VNameSource
src', Bool
b, Certs
_)) = forall s a. State s a -> s -> (a, s)
runState (forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT
(SimpleOps rep, Env rep) (State (VNameSource, Bool, Certs)) a
m (SimpleOps rep
simpl, Env rep
env)) (VNameSource
src, Bool
False, forall a. Monoid a => a
mempty)
in ((a
x, Bool
b), VNameSource
src')
askEngineEnv :: SimpleM rep (Env rep)
askEngineEnv :: forall {k} (rep :: k). SimpleM rep (Env rep)
askEngineEnv = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a, b) -> b
snd
asksEngineEnv :: (Env rep -> a) -> SimpleM rep a
asksEngineEnv :: forall {k} (rep :: k) a. (Env rep -> a) -> SimpleM rep a
asksEngineEnv Env rep -> a
f = Env rep -> a
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (rep :: k). SimpleM rep (Env rep)
askEngineEnv
askVtable :: SimpleM rep (ST.SymbolTable (Wise rep))
askVtable :: forall {k} (rep :: k). SimpleM rep (SymbolTable (Wise rep))
askVtable = forall {k} (rep :: k) a. (Env rep -> a) -> SimpleM rep a
asksEngineEnv forall {k} (rep :: k). Env rep -> SymbolTable (Wise rep)
envVtable
localVtable ::
(ST.SymbolTable (Wise rep) -> ST.SymbolTable (Wise rep)) ->
SimpleM rep a ->
SimpleM rep a
localVtable :: forall {k} (rep :: k) a.
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep a -> SimpleM rep a
localVtable SymbolTable (Wise rep) -> SymbolTable (Wise rep)
f = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \(SimpleOps rep
ops, Env rep
env) -> (SimpleOps rep
ops, Env rep
env {envVtable :: SymbolTable (Wise rep)
envVtable = SymbolTable (Wise rep) -> SymbolTable (Wise rep)
f forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Env rep -> SymbolTable (Wise rep)
envVtable Env rep
env})
collectCerts :: SimpleM rep a -> SimpleM rep (a, Certs)
collectCerts :: forall {k} (rep :: k) a. SimpleM rep a -> SimpleM rep (a, Certs)
collectCerts SimpleM rep a
m = do
a
x <- SimpleM rep a
m
(VNameSource
a, Bool
b, Certs
cs) <- forall s (m :: * -> *). MonadState s m => m s
get
forall s (m :: * -> *). MonadState s m => s -> m ()
put (VNameSource
a, Bool
b, forall a. Monoid a => a
mempty)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x, Certs
cs)
changed :: SimpleM rep ()
changed :: forall {k} (rep :: k). SimpleM rep ()
changed = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \(VNameSource
src, Bool
_, Certs
cs) -> (VNameSource
src, Bool
True, Certs
cs)
usedCerts :: Certs -> SimpleM rep ()
usedCerts :: forall {k} (rep :: k). Certs -> SimpleM rep ()
usedCerts Certs
cs = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \(VNameSource
a, Bool
b, Certs
c) -> (VNameSource
a, Bool
b, Certs
cs forall a. Semigroup a => a -> a -> a
<> Certs
c)
enterLoop :: SimpleM rep a -> SimpleM rep a
enterLoop :: forall {k} (rep :: k) a. SimpleM rep a -> SimpleM rep a
enterLoop = forall {k} (rep :: k) a.
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep a -> SimpleM rep a
localVtable forall {k} (rep :: k). SymbolTable rep -> SymbolTable rep
ST.deepen
bindFParams :: SimplifiableRep rep => [FParam (Wise rep)] -> SimpleM rep a -> SimpleM rep a
bindFParams :: forall {k} (rep :: k) a.
SimplifiableRep rep =>
[FParam (Wise rep)] -> SimpleM rep a -> SimpleM rep a
bindFParams [FParam (Wise rep)]
params =
forall {k} (rep :: k) a.
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep a -> SimpleM rep a
localVtable forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
ASTRep rep =>
[FParam rep] -> SymbolTable rep -> SymbolTable rep
ST.insertFParams [FParam (Wise rep)]
params
bindLParams :: SimplifiableRep rep => [LParam (Wise rep)] -> SimpleM rep a -> SimpleM rep a
bindLParams :: forall {k} (rep :: k) a.
SimplifiableRep rep =>
[LParam (Wise rep)] -> SimpleM rep a -> SimpleM rep a
bindLParams [LParam (Wise rep)]
params =
forall {k} (rep :: k) a.
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep a -> SimpleM rep a
localVtable forall a b. (a -> b) -> a -> b
$ \SymbolTable (Wise rep)
vtable -> forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall {k} (rep :: k).
ASTRep rep =>
LParam rep -> SymbolTable rep -> SymbolTable rep
ST.insertLParam SymbolTable (Wise rep)
vtable [LParam (Wise rep)]
params
bindArrayLParams ::
SimplifiableRep rep =>
[LParam (Wise rep)] ->
SimpleM rep a ->
SimpleM rep a
bindArrayLParams :: forall {k} (rep :: k) a.
SimplifiableRep rep =>
[LParam (Wise rep)] -> SimpleM rep a -> SimpleM rep a
bindArrayLParams [LParam (Wise rep)]
params =
forall {k} (rep :: k) a.
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep a -> SimpleM rep a
localVtable forall a b. (a -> b) -> a -> b
$ \SymbolTable (Wise 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 {k} (rep :: k).
ASTRep rep =>
LParam rep -> SymbolTable rep -> SymbolTable rep
ST.insertLParam) SymbolTable (Wise rep)
vtable [LParam (Wise rep)]
params
bindMerge ::
SimplifiableRep rep =>
[(FParam (Wise rep), SubExp, SubExpRes)] ->
SimpleM rep a ->
SimpleM rep a
bindMerge :: forall {k} (rep :: k) a.
SimplifiableRep rep =>
[(FParam (Wise rep), SubExp, SubExpRes)]
-> SimpleM rep a -> SimpleM rep a
bindMerge = forall {k} (rep :: k) a.
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep a -> SimpleM rep a
localVtable forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k).
ASTRep rep =>
[(FParam rep, SubExp, SubExpRes)]
-> SymbolTable rep -> SymbolTable rep
ST.insertLoopMerge
bindLoopVar :: SimplifiableRep rep => VName -> IntType -> SubExp -> SimpleM rep a -> SimpleM rep a
bindLoopVar :: forall {k} (rep :: k) a.
SimplifiableRep rep =>
VName -> IntType -> SubExp -> SimpleM rep a -> SimpleM rep a
bindLoopVar VName
var IntType
it SubExp
bound =
forall {k} (rep :: k) a.
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep a -> SimpleM rep a
localVtable forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
ASTRep rep =>
VName -> IntType -> SubExp -> SymbolTable rep -> SymbolTable rep
ST.insertLoopVar VName
var IntType
it SubExp
bound
makeSafe :: Exp rep -> Maybe (Exp rep)
makeSafe :: forall {k} (rep :: k). Exp rep -> Maybe (Exp rep)
makeSafe (BasicOp (BinOp (SDiv IntType
t Safety
_) SubExp
x SubExp
y)) =
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp (BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Safety -> BinOp
SDiv IntType
t Safety
Safe) SubExp
x SubExp
y)
makeSafe (BasicOp (BinOp (SDivUp IntType
t Safety
_) SubExp
x SubExp
y)) =
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp (BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Safety -> BinOp
SDivUp IntType
t Safety
Safe) SubExp
x SubExp
y)
makeSafe (BasicOp (BinOp (SQuot IntType
t Safety
_) SubExp
x SubExp
y)) =
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp (BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Safety -> BinOp
SQuot IntType
t Safety
Safe) SubExp
x SubExp
y)
makeSafe (BasicOp (BinOp (UDiv IntType
t Safety
_) SubExp
x SubExp
y)) =
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp (BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Safety -> BinOp
UDiv IntType
t Safety
Safe) SubExp
x SubExp
y)
makeSafe (BasicOp (BinOp (UDivUp IntType
t Safety
_) SubExp
x SubExp
y)) =
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp (BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Safety -> BinOp
UDivUp IntType
t Safety
Safe) SubExp
x SubExp
y)
makeSafe (BasicOp (BinOp (SMod IntType
t Safety
_) SubExp
x SubExp
y)) =
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp (BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Safety -> BinOp
SMod IntType
t Safety
Safe) SubExp
x SubExp
y)
makeSafe (BasicOp (BinOp (SRem IntType
t Safety
_) SubExp
x SubExp
y)) =
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp (BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Safety -> BinOp
SRem IntType
t Safety
Safe) SubExp
x SubExp
y)
makeSafe (BasicOp (BinOp (UMod IntType
t Safety
_) SubExp
x SubExp
y)) =
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp (BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Safety -> BinOp
UMod IntType
t Safety
Safe) SubExp
x SubExp
y)
makeSafe Exp rep
_ =
forall a. Maybe a
Nothing
emptyOfType :: MonadBuilder m => [VName] -> Type -> m (Exp (Rep m))
emptyOfType :: forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Type -> m (Exp (Rep m))
emptyOfType [VName]
_ Mem {} =
forall a. HasCallStack => [Char] -> a
error [Char]
"emptyOfType: Cannot hoist non-existential memory."
emptyOfType [VName]
_ Acc {} =
forall a. HasCallStack => [Char] -> a
error [Char]
"emptyOfType: Cannot hoist accumulator."
emptyOfType [VName]
_ (Prim PrimType
pt) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant forall a b. (a -> b) -> a -> b
$ PrimType -> PrimValue
blankPrimValue PrimType
pt
emptyOfType [VName]
ctx_names (Array PrimType
et Shape
shape NoUniqueness
_) = do
let dims :: [SubExp]
dims = forall a b. (a -> b) -> [a] -> [b]
map SubExp -> SubExp
zeroIfContext forall a b. (a -> b) -> a -> b
$ forall d. ShapeBase d -> [d]
shapeDims Shape
shape
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ PrimType -> [SubExp] -> BasicOp
Scratch PrimType
et [SubExp]
dims
where
zeroIfContext :: SubExp -> SubExp
zeroIfContext (Var VName
v) | VName
v forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
ctx_names = IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
zeroIfContext SubExp
se = SubExp
se
protectIf ::
MonadBuilder m =>
Protect m ->
(Exp (Rep m) -> Bool) ->
SubExp ->
Stm (Rep m) ->
m ()
protectIf :: forall (m :: * -> *).
MonadBuilder m =>
Protect m -> (Exp (Rep m) -> Bool) -> SubExp -> Stm (Rep m) -> m ()
protectIf Protect m
_ Exp (Rep m) -> Bool
_ SubExp
taken (Let Pat (LetDec (Rep m))
pat StmAux (ExpDec (Rep m))
aux (Match [SubExp
cond] [Case [Just (BoolValue Bool
True)] Body (Rep m)
taken_body] Body (Rep m)
untaken_body (MatchDec [BranchType (Rep m)]
if_ts MatchSort
MatchFallback))) = do
SubExp
cond' <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"protect_cond_conj" forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp BinOp
LogAnd SubExp
taken SubExp
cond
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux (ExpDec (Rep m))
aux forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
MonadBuilder m =>
Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
letBind Pat (LetDec (Rep m))
pat forall a b. (a -> b) -> a -> b
$
forall {k} (rep :: k).
[SubExp]
-> [Case (Body rep)]
-> Body rep
-> MatchDec (BranchType rep)
-> Exp rep
Match [SubExp
cond'] [forall body. [Maybe PrimValue] -> body -> Case body
Case [forall a. a -> Maybe a
Just (Bool -> PrimValue
BoolValue Bool
True)] Body (Rep m)
taken_body] Body (Rep m)
untaken_body forall a b. (a -> b) -> a -> b
$
forall rt. [rt] -> MatchSort -> MatchDec rt
MatchDec [BranchType (Rep m)]
if_ts MatchSort
MatchFallback
protectIf Protect m
_ Exp (Rep m) -> Bool
_ SubExp
taken (Let Pat (LetDec (Rep m))
pat StmAux (ExpDec (Rep m))
aux (BasicOp (Assert SubExp
cond ErrorMsg SubExp
msg (SrcLoc, [SrcLoc])
loc))) = do
SubExp
not_taken <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"loop_not_taken" forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
UnOp UnOp
Not SubExp
taken
SubExp
cond' <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"protect_assert_disj" forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp BinOp
LogOr SubExp
not_taken SubExp
cond
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux (ExpDec (Rep m))
aux forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
letBind Pat (LetDec (Rep m))
pat forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> ErrorMsg SubExp -> (SrcLoc, [SrcLoc]) -> BasicOp
Assert SubExp
cond' ErrorMsg SubExp
msg (SrcLoc, [SrcLoc])
loc
protectIf Protect m
protect Exp (Rep m) -> Bool
_ SubExp
taken (Let Pat (LetDec (Rep m))
pat StmAux (ExpDec (Rep m))
aux (Op Op (Rep m)
op))
| Just m ()
m <- Protect m
protect SubExp
taken Pat (LetDec (Rep m))
pat Op (Rep m)
op =
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux (ExpDec (Rep m))
aux m ()
m
protectIf Protect m
_ Exp (Rep m) -> Bool
f SubExp
taken (Let Pat (LetDec (Rep m))
pat StmAux (ExpDec (Rep m))
aux Exp (Rep m)
e)
| Exp (Rep m) -> Bool
f Exp (Rep m)
e =
case forall {k} (rep :: k). Exp rep -> Maybe (Exp rep)
makeSafe Exp (Rep m)
e of
Just Exp (Rep m)
e' ->
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux (ExpDec (Rep m))
aux forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
letBind Pat (LetDec (Rep m))
pat Exp (Rep m)
e'
Maybe (Exp (Rep m))
Nothing -> do
Body (Rep m)
taken_body <- forall (m :: * -> *).
MonadBuilder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody [forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp (Rep m)
e]
Body (Rep m)
untaken_body <-
forall (m :: * -> *).
MonadBuilder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Type -> m (Exp (Rep m))
emptyOfType forall a b. (a -> b) -> a -> b
$ forall dec. Pat dec -> [VName]
patNames Pat (LetDec (Rep m))
pat) (forall dec. Typed dec => Pat dec -> [Type]
patTypes Pat (LetDec (Rep m))
pat)
[BranchType (Rep m)]
if_ts <- forall {k} (rep :: k) (m :: * -> *).
(ASTRep rep, HasScope rep m, Monad m) =>
Pat (LetDec rep) -> m [BranchType rep]
expTypesFromPat Pat (LetDec (Rep m))
pat
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux (ExpDec (Rep m))
aux forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
MonadBuilder m =>
Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
letBind Pat (LetDec (Rep m))
pat
forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
[SubExp]
-> [Case (Body rep)]
-> Body rep
-> MatchDec (BranchType rep)
-> Exp rep
Match
[SubExp
taken]
[forall body. [Maybe PrimValue] -> body -> Case body
Case [forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue Bool
True] Body (Rep m)
taken_body]
Body (Rep m)
untaken_body
forall a b. (a -> b) -> a -> b
$ forall rt. [rt] -> MatchSort -> MatchDec rt
MatchDec [BranchType (Rep m)]
if_ts MatchSort
MatchFallback
protectIf Protect m
_ Exp (Rep m) -> Bool
_ SubExp
_ Stm (Rep m)
stm =
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm Stm (Rep m)
stm
protectLoopHoisted ::
SimplifiableRep rep =>
[(FParam (Wise rep), SubExp)] ->
LoopForm (Wise rep) ->
SimpleM rep (a, b, Stms (Wise rep)) ->
SimpleM rep (a, b, Stms (Wise rep))
protectLoopHoisted :: forall {k} (rep :: k) a b.
SimplifiableRep rep =>
[(FParam (Wise rep), SubExp)]
-> LoopForm (Wise rep)
-> SimpleM rep (a, b, Stms (Wise rep))
-> SimpleM rep (a, b, Stms (Wise rep))
protectLoopHoisted [(FParam (Wise rep), SubExp)]
merge LoopForm (Wise rep)
form SimpleM rep (a, b, Stms (Wise rep))
m = do
(a
x, b
y, Stms (Wise rep)
stms) <- SimpleM rep (a, b, Stms (Wise rep))
m
SubExp
-> Pat (VarWisdom, LetDec rep)
-> OpWithWisdom (Op rep)
-> Maybe (Builder (Wise rep) ())
ops <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
SimpleOps rep -> Protect (Builder (Wise rep))
protectHoistedOpS forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst
Stms (Wise rep)
stms' <- forall {k1} {k2} (m :: * -> *) (somerep :: k1) (rep :: k2) a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (Stms rep)
runBuilder_ forall a b. (a -> b) -> a -> b
$ do
if Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall {k} (rep :: k). IsOp (Op rep) => Exp rep -> Bool
safeExp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). Stm rep -> Exp rep
stmExp) Stms (Wise rep)
stms
then do
SubExp
is_nonempty <- BuilderT (Wise rep) (State VNameSource) SubExp
checkIfNonEmpty
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall (m :: * -> *).
MonadBuilder m =>
Protect m -> (Exp (Rep m) -> Bool) -> SubExp -> Stm (Rep m) -> m ()
protectIf SubExp
-> Pat (VarWisdom, LetDec rep)
-> OpWithWisdom (Op rep)
-> Maybe (Builder (Wise rep) ())
ops (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). IsOp (Op rep) => Exp rep -> Bool
safeExp) SubExp
is_nonempty) Stms (Wise rep)
stms
else forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms Stms (Wise rep)
stms
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x, b
y, Stms (Wise rep)
stms')
where
checkIfNonEmpty :: BuilderT (Wise rep) (State VNameSource) SubExp
checkIfNonEmpty =
case LoopForm (Wise rep)
form of
WhileLoop VName
cond
| Just (Param (FParamInfo rep)
_, SubExp
cond_init) <-
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall a. Eq a => a -> a -> Bool
== VName
cond) 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 (Wise rep), SubExp)]
merge ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
cond_init
| Bool
otherwise -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall v. IsValue v => v -> SubExp
constant Bool
True
ForLoop VName
_ IntType
it SubExp
bound [(LParam (Wise rep), VName)]
_ ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"loop_nonempty" forall a b. (a -> b) -> a -> b
$
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (IntType -> CmpOp
CmpSlt IntType
it) (IntType -> Integer -> SubExp
intConst IntType
it Integer
0) SubExp
bound
matching ::
BuilderOps rep =>
[(SubExp, Maybe PrimValue)] ->
Builder rep SubExp
matching :: forall rep.
BuilderOps rep =>
[(SubExp, Maybe PrimValue)] -> Builder rep SubExp
matching = forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"match" forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe forall {f :: * -> *}.
MonadBuilder f =>
(SubExp, Maybe PrimValue) -> Maybe (f SubExp)
cmp
where
cmp :: (SubExp, Maybe PrimValue) -> Maybe (f SubExp)
cmp (SubExp
se, Just (BoolValue Bool
True)) =
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se
cmp (SubExp
se, Just PrimValue
v) =
forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"match_val" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (PrimType -> CmpOp
CmpEq (PrimValue -> PrimType
primValueType PrimValue
v)) SubExp
se (PrimValue -> SubExp
Constant PrimValue
v)
cmp (SubExp
_, Maybe PrimValue
Nothing) = forall a. Maybe a
Nothing
matchingExactlyThis ::
BuilderOps rep =>
[SubExp] ->
[[Maybe PrimValue]] ->
[Maybe PrimValue] ->
Builder rep SubExp
matchingExactlyThis :: forall rep.
BuilderOps rep =>
[SubExp]
-> [[Maybe PrimValue]] -> [Maybe PrimValue] -> Builder rep SubExp
matchingExactlyThis [SubExp]
ses [[Maybe PrimValue]]
prior [Maybe PrimValue]
this = do
[SubExp]
prior_matches <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall rep.
BuilderOps rep =>
[(SubExp, Maybe PrimValue)] -> Builder rep SubExp
matching forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
ses) [[Maybe PrimValue]]
prior
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"matching_just_this"
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp
BinOp
LogAnd
(forall (m :: * -> *).
MonadBuilder m =>
UnOp -> m (Exp (Rep m)) -> m (Exp (Rep m))
eUnOp UnOp
Not (forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAny [SubExp]
prior_matches))
(forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall rep.
BuilderOps rep =>
[(SubExp, Maybe PrimValue)] -> Builder rep SubExp
matching (forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
ses [Maybe PrimValue]
this))
protectCaseHoisted ::
SimplifiableRep rep =>
[SubExp] ->
[[Maybe PrimValue]] ->
[Maybe PrimValue] ->
SimpleM rep (Stms (Wise rep), a) ->
SimpleM rep (Stms (Wise rep), a)
protectCaseHoisted :: forall {k} (rep :: k) a.
SimplifiableRep rep =>
[SubExp]
-> [[Maybe PrimValue]]
-> [Maybe PrimValue]
-> SimpleM rep (Stms (Wise rep), a)
-> SimpleM rep (Stms (Wise rep), a)
protectCaseHoisted [SubExp]
ses [[Maybe PrimValue]]
prior [Maybe PrimValue]
vs SimpleM rep (Stms (Wise rep), a)
m = do
(Stms (Wise rep)
hoisted, a
x) <- SimpleM rep (Stms (Wise rep), a)
m
SubExp
-> Pat (VarWisdom, LetDec rep)
-> OpWithWisdom (Op rep)
-> Maybe (Builder (Wise rep) ())
ops <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
SimpleOps rep -> Protect (Builder (Wise rep))
protectHoistedOpS forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst
Stms (Wise rep)
hoisted' <- forall {k1} {k2} (m :: * -> *) (somerep :: k1) (rep :: k2) a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (Stms rep)
runBuilder_ forall a b. (a -> b) -> a -> b
$ do
if Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall {k} (rep :: k). IsOp (Op rep) => Exp rep -> Bool
safeExp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). Stm rep -> Exp rep
stmExp) Stms (Wise rep)
hoisted
then do
SubExp
cond' <- forall rep.
BuilderOps rep =>
[SubExp]
-> [[Maybe PrimValue]] -> [Maybe PrimValue] -> Builder rep SubExp
matchingExactlyThis [SubExp]
ses [[Maybe PrimValue]]
prior [Maybe PrimValue]
vs
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall (m :: * -> *).
MonadBuilder m =>
Protect m -> (Exp (Rep m) -> Bool) -> SubExp -> Stm (Rep m) -> m ()
protectIf SubExp
-> Pat (VarWisdom, LetDec rep)
-> OpWithWisdom (Op rep)
-> Maybe (Builder (Wise rep) ())
ops forall {k} {rep :: k}. ASTRep rep => Exp rep -> Bool
unsafeOrCostly SubExp
cond') Stms (Wise rep)
hoisted
else forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms Stms (Wise rep)
hoisted
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms (Wise rep)
hoisted', a
x)
where
unsafeOrCostly :: Exp rep -> Bool
unsafeOrCostly Exp rep
e = Bool -> Bool
not (forall {k} (rep :: k). IsOp (Op rep) => Exp rep -> Bool
safeExp Exp rep
e) Bool -> Bool -> Bool
|| Bool -> Bool
not (forall {k} {rep :: k}. ASTRep rep => Exp rep -> Bool
cheapExp Exp rep
e)
notWorthHoisting :: ASTRep rep => BlockPred rep
notWorthHoisting :: forall {k} (rep :: k). ASTRep rep => BlockPred rep
notWorthHoisting SymbolTable rep
_ UsageTable
_ (Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
_ Exp rep
e) =
Bool -> Bool
not (forall {k} (rep :: k). IsOp (Op rep) => Exp rep -> Bool
safeExp Exp rep
e) Bool -> Bool -> Bool
&& forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((forall a. Ord a => a -> a -> Bool
> Int
0) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall shape u. ArrayShape shape => TypeBase shape u -> Int
arrayRank) (forall dec. Typed dec => Pat dec -> [Type]
patTypes Pat (LetDec rep)
pat)
nonrecSimplifyStm ::
SimplifiableRep rep =>
Stm (Wise rep) ->
SimpleM rep (Stm (Wise rep))
nonrecSimplifyStm :: forall {k} (rep :: k).
SimplifiableRep rep =>
Stm (Wise rep) -> SimpleM rep (Stm (Wise rep))
nonrecSimplifyStm (Let Pat (LetDec (Wise rep))
pat (StmAux Certs
cs Attrs
attrs (ExpWisdom
_, ExpDec rep
dec)) Exp (Wise rep)
e) = do
Certs
cs' <- forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify Certs
cs
Exp (Wise rep)
e' <- forall {k} (rep :: k).
SimplifiableRep rep =>
Exp (Wise rep) -> SimpleM rep (Exp (Wise rep))
simplifyExpBase Exp (Wise rep)
e
Pat (LetDec rep)
-> Exp (Wise rep) -> SimpleM rep (Pat (LetDec rep))
simplifyPat <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
SimpleOps rep
-> Pat (LetDec rep)
-> Exp (Wise rep)
-> SimpleM rep (Pat (LetDec rep))
simplifyPatFromExpS forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst
(Pat (LetDec rep)
pat', Certs
pat_cs) <- forall {k} (rep :: k) a. SimpleM rep a -> SimpleM rep (a, Certs)
collectCerts forall a b. (a -> b) -> a -> b
$ Pat (LetDec rep)
-> Exp (Wise rep) -> SimpleM rep (Pat (LetDec rep))
simplifyPat (forall a. Pat (VarWisdom, a) -> Pat a
removePatWisdom Pat (LetDec (Wise rep))
pat) Exp (Wise rep)
e'
let aux' :: StmAux (ExpDec rep)
aux' = forall dec. Certs -> Attrs -> dec -> StmAux dec
StmAux (Certs
cs' forall a. Semigroup a => a -> a -> a
<> Certs
pat_cs) Attrs
attrs ExpDec rep
dec
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
(ASTRep rep, CanBeWise (Op rep)) =>
Pat (LetDec rep)
-> StmAux (ExpDec rep) -> Exp (Wise rep) -> Stm (Wise rep)
mkWiseStm Pat (LetDec rep)
pat' StmAux (ExpDec rep)
aux' Exp (Wise rep)
e'
recSimplifyStm ::
SimplifiableRep rep =>
Stm (Wise rep) ->
UT.UsageTable ->
SimpleM rep (Stms (Wise rep), Stm (Wise rep))
recSimplifyStm :: forall {k} (rep :: k).
SimplifiableRep rep =>
Stm (Wise rep)
-> UsageTable -> SimpleM rep (Stms (Wise rep), Stm (Wise rep))
recSimplifyStm (Let Pat (LetDec (Wise rep))
pat (StmAux Certs
cs Attrs
attrs (ExpWisdom
_, ExpDec rep
dec)) Exp (Wise rep)
e) UsageTable
usage = do
((Exp (Wise rep)
e', Stms (Wise rep)
e_hoisted), Certs
e_cs) <- forall {k} (rep :: k) a. SimpleM rep a -> SimpleM rep (a, Certs)
collectCerts forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
SimplifiableRep rep =>
UsageTable
-> Pat (LetDec (Wise rep))
-> Exp (Wise rep)
-> SimpleM rep (Exp (Wise rep), Stms (Wise rep))
simplifyExp UsageTable
usage Pat (LetDec (Wise rep))
pat Exp (Wise rep)
e
let aux' :: StmAux (ExpDec rep)
aux' = forall dec. Certs -> Attrs -> dec -> StmAux dec
StmAux (Certs
cs forall a. Semigroup a => a -> a -> a
<> Certs
e_cs) Attrs
attrs ExpDec rep
dec
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms (Wise rep)
e_hoisted, forall {k} (rep :: k).
(ASTRep rep, CanBeWise (Op rep)) =>
Pat (LetDec rep)
-> StmAux (ExpDec rep) -> Exp (Wise rep) -> Stm (Wise rep)
mkWiseStm (forall a. Pat (VarWisdom, a) -> Pat a
removePatWisdom Pat (LetDec (Wise rep))
pat) StmAux (ExpDec rep)
aux' Exp (Wise rep)
e')
hoistStms ::
SimplifiableRep rep =>
RuleBook (Wise rep) ->
BlockPred (Wise rep) ->
Stms (Wise rep) ->
SimpleM rep (a, UT.UsageTable) ->
SimpleM rep (a, Stms (Wise rep), Stms (Wise rep))
hoistStms :: forall {k} (rep :: k) a.
SimplifiableRep rep =>
RuleBook (Wise rep)
-> BlockPred (Wise rep)
-> Stms (Wise rep)
-> SimpleM rep (a, UsageTable)
-> SimpleM rep (a, Stms (Wise rep), Stms (Wise rep))
hoistStms RuleBook (Wise rep)
rules BlockPred (Wise rep)
block Stms (Wise rep)
orig_stms SimpleM rep (a, UsageTable)
final = do
(a
a, [Stm (Wise rep)]
blocked, [Stm (Wise rep)]
hoisted) <- Stms (Wise rep)
-> SimpleM rep (a, [Stm (Wise rep)], [Stm (Wise rep)])
simplifyStmsBottomUp Stms (Wise rep)
orig_stms
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Stm (Wise rep)]
hoisted) forall {k} (rep :: k). SimpleM rep ()
changed
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, forall {k} (rep :: k). [Stm rep] -> Stms rep
stmsFromList [Stm (Wise rep)]
blocked, forall {k} (rep :: k). [Stm rep] -> Stms rep
stmsFromList [Stm (Wise rep)]
hoisted)
where
simplifyStmsBottomUp :: Stms (Wise rep)
-> SimpleM rep (a, [Stm (Wise rep)], [Stm (Wise rep)])
simplifyStmsBottomUp Stms (Wise rep)
stms = do
OpWithWisdom (Op rep) -> UsageTable
opUsage <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). SimpleOps rep -> Op (Wise rep) -> UsageTable
opUsageS forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst
let usageInStm :: Stm (Wise rep) -> UsageTable
usageInStm Stm (Wise rep)
stm =
forall {k} (rep :: k).
(ASTRep rep, Aliased rep) =>
Stm rep -> UsageTable
UT.usageInStm Stm (Wise rep)
stm
forall a. Semigroup a => a -> a -> a
<> case forall {k} (rep :: k). Stm rep -> Exp rep
stmExp Stm (Wise rep)
stm of
Op Op (Wise rep)
op -> OpWithWisdom (Op rep) -> UsageTable
opUsage Op (Wise rep)
op
Exp (Wise rep)
_ -> forall a. Monoid a => a
mempty
(a
x, UsageTable
_, [Either (Stm (Wise rep)) (Stm (Wise rep))]
stms') <- (Stm (Wise rep) -> UsageTable)
-> Stms (Wise rep)
-> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
hoistableStms Stm (Wise rep) -> UsageTable
usageInStm Stms (Wise rep)
stms
let ([Stm (Wise rep)]
blocked, [Stm (Wise rep)]
hoisted) = forall a b. [Either a b] -> ([a], [b])
partitionEithers forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
ASTRep rep =>
[Either (Stm rep) (Stm rep)] -> [Either (Stm rep) (Stm rep)]
blockUnhoistedDeps [Either (Stm (Wise rep)) (Stm (Wise rep))]
stms'
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x, [Stm (Wise rep)]
blocked, [Stm (Wise rep)]
hoisted)
descend :: (Stm (Wise rep) -> UsageTable)
-> Stms (Wise rep)
-> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
-> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
descend Stm (Wise rep) -> UsageTable
usageInStm Stms (Wise rep)
stms SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
m =
case forall {k} (rep :: k). Stms rep -> Maybe (Stm rep, Stms rep)
stmsHead Stms (Wise rep)
stms of
Maybe (Stm (Wise rep), Stms (Wise rep))
Nothing -> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
m
Just (Stm (Wise rep)
stms_h, Stms (Wise rep)
stms_t) -> forall {k} (rep :: k) a.
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep a -> SimpleM rep a
localVtable (forall {k} (rep :: k).
(ASTRep rep, IndexOp (Op rep), Aliased rep) =>
Stm rep -> SymbolTable rep -> SymbolTable rep
ST.insertStm Stm (Wise rep)
stms_h) forall a b. (a -> b) -> a -> b
$ do
(a
x, UsageTable
usage, [Either (Stm (Wise rep)) (Stm (Wise rep))]
stms_t') <- (Stm (Wise rep) -> UsageTable)
-> Stms (Wise rep)
-> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
-> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
descend Stm (Wise rep) -> UsageTable
usageInStm Stms (Wise rep)
stms_t SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
m
(Stm (Wise rep) -> UsageTable)
-> Stm (Wise rep)
-> [Either (Stm (Wise rep)) (Stm (Wise rep))]
-> UsageTable
-> a
-> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
process Stm (Wise rep) -> UsageTable
usageInStm Stm (Wise rep)
stms_h [Either (Stm (Wise rep)) (Stm (Wise rep))]
stms_t' UsageTable
usage a
x
process :: (Stm (Wise rep) -> UsageTable)
-> Stm (Wise rep)
-> [Either (Stm (Wise rep)) (Stm (Wise rep))]
-> UsageTable
-> a
-> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
process Stm (Wise rep) -> UsageTable
usageInStm Stm (Wise rep)
stm [Either (Stm (Wise rep)) (Stm (Wise rep))]
stms UsageTable
usage a
x = do
SymbolTable (Wise rep)
vtable <- forall {k} (rep :: k). SimpleM rep (SymbolTable (Wise rep))
askVtable
Maybe (Stms (Wise rep))
res <- forall {k} (m :: * -> *) (rep :: k).
(MonadFreshNames m, HasScope rep m) =>
RuleBook rep
-> (SymbolTable rep, UsageTable) -> Stm rep -> m (Maybe (Stms rep))
bottomUpSimplifyStm RuleBook (Wise rep)
rules (SymbolTable (Wise rep)
vtable, UsageTable
usage) Stm (Wise rep)
stm
case Maybe (Stms (Wise rep))
res of
Maybe (Stms (Wise rep))
Nothing
| BlockPred (Wise rep)
block SymbolTable (Wise rep)
vtable UsageTable
usage Stm (Wise rep)
stm ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( a
x,
forall {k} (rep :: k).
(ASTRep rep, Aliased rep) =>
(Stm rep -> UsageTable)
-> SymbolTable rep -> UsageTable -> Stm rep -> UsageTable
expandUsage Stm (Wise rep) -> UsageTable
usageInStm SymbolTable (Wise rep)
vtable UsageTable
usage Stm (Wise rep)
stm
UsageTable -> [VName] -> UsageTable
`UT.without` forall {k} (rep :: k). Stm rep -> [VName]
provides Stm (Wise rep)
stm,
forall a b. a -> Either a b
Left Stm (Wise rep)
stm forall a. a -> [a] -> [a]
: [Either (Stm (Wise rep)) (Stm (Wise rep))]
stms
)
| Bool
otherwise ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( a
x,
forall {k} (rep :: k).
(ASTRep rep, Aliased rep) =>
(Stm rep -> UsageTable)
-> SymbolTable rep -> UsageTable -> Stm rep -> UsageTable
expandUsage Stm (Wise rep) -> UsageTable
usageInStm SymbolTable (Wise rep)
vtable UsageTable
usage Stm (Wise rep)
stm,
forall a b. b -> Either a b
Right Stm (Wise rep)
stm forall a. a -> [a] -> [a]
: [Either (Stm (Wise rep)) (Stm (Wise rep))]
stms
)
Just Stms (Wise rep)
optimstms -> do
forall {k} (rep :: k). SimpleM rep ()
changed
(Stm (Wise rep) -> UsageTable)
-> Stms (Wise rep)
-> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
-> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
descend Stm (Wise rep) -> UsageTable
usageInStm Stms (Wise rep)
optimstms forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x, UsageTable
usage, [Either (Stm (Wise rep)) (Stm (Wise rep))]
stms)
hoistableStms :: (Stm (Wise rep) -> UsageTable)
-> Stms (Wise rep)
-> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
hoistableStms Stm (Wise rep) -> UsageTable
usageInStm Stms (Wise rep)
stms =
case forall {k} (rep :: k). Stms rep -> Maybe (Stm rep, Stms rep)
stmsHead Stms (Wise rep)
stms of
Maybe (Stm (Wise rep), Stms (Wise rep))
Nothing -> do
(a
x, UsageTable
usage) <- SimpleM rep (a, UsageTable)
final
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x, UsageTable
usage, forall a. Monoid a => a
mempty)
Just (Stm (Wise rep)
stms_h, Stms (Wise rep)
stms_t) -> do
Stm (Wise rep)
stms_h' <- forall {k} (rep :: k).
SimplifiableRep rep =>
Stm (Wise rep) -> SimpleM rep (Stm (Wise rep))
nonrecSimplifyStm Stm (Wise rep)
stms_h
SymbolTable (Wise rep)
vtable <- forall {k} (rep :: k). SimpleM rep (SymbolTable (Wise rep))
askVtable
Maybe (Stms (Wise rep))
simplified <- forall {k} (m :: * -> *) (rep :: k).
(MonadFreshNames m, HasScope rep m) =>
RuleBook rep -> SymbolTable rep -> Stm rep -> m (Maybe (Stms rep))
topDownSimplifyStm RuleBook (Wise rep)
rules SymbolTable (Wise rep)
vtable Stm (Wise rep)
stms_h'
case Maybe (Stms (Wise rep))
simplified of
Just Stms (Wise rep)
newstms -> do
forall {k} (rep :: k). SimpleM rep ()
changed
(Stm (Wise rep) -> UsageTable)
-> Stms (Wise rep)
-> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
hoistableStms Stm (Wise rep) -> UsageTable
usageInStm (Stms (Wise rep)
newstms forall a. Semigroup a => a -> a -> a
<> Stms (Wise rep)
stms_t)
Maybe (Stms (Wise rep))
Nothing -> do
(a
x, UsageTable
usage, [Either (Stm (Wise rep)) (Stm (Wise rep))]
stms_t') <-
forall {k} (rep :: k) a.
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep a -> SimpleM rep a
localVtable (forall {k} (rep :: k).
(ASTRep rep, IndexOp (Op rep), Aliased rep) =>
Stm rep -> SymbolTable rep -> SymbolTable rep
ST.insertStm Stm (Wise rep)
stms_h') forall a b. (a -> b) -> a -> b
$
(Stm (Wise rep) -> UsageTable)
-> Stms (Wise rep)
-> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
hoistableStms Stm (Wise rep) -> UsageTable
usageInStm Stms (Wise rep)
stms_t
if Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> UsageTable -> Bool
`UT.isUsedDirectly` UsageTable
usage) forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Stm rep -> [VName]
provides Stm (Wise rep)
stms_h'
then
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x, UsageTable
usage, [Either (Stm (Wise rep)) (Stm (Wise rep))]
stms_t')
else do
(Stms (Wise rep)
stms_h_stms, Stm (Wise rep)
stms_h'') <- forall {k} (rep :: k).
SimplifiableRep rep =>
Stm (Wise rep)
-> UsageTable -> SimpleM rep (Stms (Wise rep), Stm (Wise rep))
recSimplifyStm Stm (Wise rep)
stms_h' UsageTable
usage
(Stm (Wise rep) -> UsageTable)
-> Stms (Wise rep)
-> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
-> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
descend Stm (Wise rep) -> UsageTable
usageInStm Stms (Wise rep)
stms_h_stms forall a b. (a -> b) -> a -> b
$
(Stm (Wise rep) -> UsageTable)
-> Stm (Wise rep)
-> [Either (Stm (Wise rep)) (Stm (Wise rep))]
-> UsageTable
-> a
-> SimpleM
rep (a, UsageTable, [Either (Stm (Wise rep)) (Stm (Wise rep))])
process Stm (Wise rep) -> UsageTable
usageInStm Stm (Wise rep)
stms_h'' [Either (Stm (Wise rep)) (Stm (Wise rep))]
stms_t' UsageTable
usage a
x
blockUnhoistedDeps ::
ASTRep rep =>
[Either (Stm rep) (Stm rep)] ->
[Either (Stm rep) (Stm rep)]
blockUnhoistedDeps :: forall {k} (rep :: k).
ASTRep rep =>
[Either (Stm rep) (Stm rep)] -> [Either (Stm rep) (Stm rep)]
blockUnhoistedDeps = forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL forall {k} {rep :: k}.
(FreeDec (ExpDec rep), FreeDec (BodyDec rep),
FreeIn (FParamInfo rep), FreeIn (LParamInfo rep),
FreeIn (LetDec rep), FreeIn (RetType rep), FreeIn (BranchType rep),
FreeIn (Op rep)) =>
Names
-> Either (Stm rep) (Stm rep)
-> (Names, Either (Stm rep) (Stm rep))
block forall a. Monoid a => a
mempty
where
block :: Names
-> Either (Stm rep) (Stm rep)
-> (Names, Either (Stm rep) (Stm rep))
block Names
blocked (Left Stm rep
need) =
(Names
blocked forall a. Semigroup a => a -> a -> a
<> [VName] -> Names
namesFromList (forall {k} (rep :: k). Stm rep -> [VName]
provides Stm rep
need), forall a b. a -> Either a b
Left Stm rep
need)
block Names
blocked (Right Stm rep
need)
| Names
blocked Names -> Names -> Bool
`namesIntersect` forall a. FreeIn a => a -> Names
freeIn Stm rep
need =
(Names
blocked forall a. Semigroup a => a -> a -> a
<> [VName] -> Names
namesFromList (forall {k} (rep :: k). Stm rep -> [VName]
provides Stm rep
need), forall a b. a -> Either a b
Left Stm rep
need)
| Bool
otherwise =
(Names
blocked, forall a b. b -> Either a b
Right Stm rep
need)
provides :: Stm rep -> [VName]
provides :: forall {k} (rep :: k). Stm rep -> [VName]
provides = forall dec. Pat dec -> [VName]
patNames forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). Stm rep -> Pat (LetDec rep)
stmPat
expandUsage ::
(ASTRep rep, Aliased rep) =>
(Stm rep -> UT.UsageTable) ->
ST.SymbolTable rep ->
UT.UsageTable ->
Stm rep ->
UT.UsageTable
expandUsage :: forall {k} (rep :: k).
(ASTRep rep, Aliased rep) =>
(Stm rep -> UsageTable)
-> SymbolTable rep -> UsageTable -> Stm rep -> UsageTable
expandUsage Stm rep -> UsageTable
usageInStm SymbolTable rep
vtable UsageTable
utable stm :: Stm rep
stm@(Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
aux Exp rep
e) =
UsageTable
stmUsages forall a. Semigroup a => a -> a -> a
<> UsageTable
utable
where
stmUsages :: UsageTable
stmUsages =
(VName -> Names) -> UsageTable -> UsageTable
UT.expand (forall {k} (rep :: k). VName -> SymbolTable rep -> Names
`ST.lookupAliases` SymbolTable rep
vtable) (Stm rep -> UsageTable
usageInStm Stm rep
stm forall a. Semigroup a => a -> a -> a
<> UsageTable
usageThroughAliases)
forall a. Semigroup a => a -> a -> a
<> ( if forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> UsageTable -> Bool
`UT.isSize` UsageTable
utable) (forall dec. Pat dec -> [VName]
patNames Pat (LetDec rep)
pat)
then Names -> UsageTable
UT.sizeUsages (forall a. FreeIn a => a -> Names
freeIn (forall dec. StmAux dec -> Certs
stmAuxCerts StmAux (ExpDec rep)
aux) forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn Exp rep
e)
else forall a. Monoid a => a
mempty
)
usageThroughAliases :: UsageTable
usageThroughAliases =
forall a. Monoid a => [a] -> a
mconcat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (VName, Names) -> Maybe UsageTable
usageThroughBindeeAliases forall a b. (a -> b) -> a -> b
$
forall a b. [a] -> [b] -> [(a, b)]
zip (forall dec. Pat dec -> [VName]
patNames Pat (LetDec rep)
pat) (forall dec. AliasesOf dec => Pat dec -> [Names]
patAliases Pat (LetDec rep)
pat)
usageThroughBindeeAliases :: (VName, Names) -> Maybe UsageTable
usageThroughBindeeAliases (VName
name, Names
aliases) = do
Usages
uses <- VName -> UsageTable -> Maybe Usages
UT.lookup VName
name UsageTable
utable
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$
forall a b. (a -> b) -> [a] -> [b]
map (VName -> Usages -> UsageTable
`UT.usage` (Usages
uses Usages -> Usages -> Usages
`UT.withoutU` Usages
UT.presentU)) forall a b. (a -> b) -> a -> b
$
Names -> [VName]
namesToList Names
aliases
type BlockPred rep = ST.SymbolTable rep -> UT.UsageTable -> Stm rep -> Bool
neverBlocks :: BlockPred rep
neverBlocks :: forall {k} (rep :: k). BlockPred rep
neverBlocks SymbolTable rep
_ UsageTable
_ Stm rep
_ = Bool
False
alwaysBlocks :: BlockPred rep
alwaysBlocks :: forall {k} (rep :: k). BlockPred rep
alwaysBlocks SymbolTable rep
_ UsageTable
_ Stm rep
_ = Bool
True
isFalse :: Bool -> BlockPred rep
isFalse :: forall {k} (rep :: k). Bool -> BlockPred rep
isFalse Bool
b SymbolTable rep
_ UsageTable
_ Stm rep
_ = Bool -> Bool
not Bool
b
orIf :: BlockPred rep -> BlockPred rep -> BlockPred rep
orIf :: forall {k} (rep :: k).
BlockPred rep -> BlockPred rep -> BlockPred rep
orIf BlockPred rep
p1 BlockPred rep
p2 SymbolTable rep
body UsageTable
vtable Stm rep
need = BlockPred rep
p1 SymbolTable rep
body UsageTable
vtable Stm rep
need Bool -> Bool -> Bool
|| BlockPred rep
p2 SymbolTable rep
body UsageTable
vtable Stm rep
need
andAlso :: BlockPred rep -> BlockPred rep -> BlockPred rep
andAlso :: forall {k} (rep :: k).
BlockPred rep -> BlockPred rep -> BlockPred rep
andAlso BlockPred rep
p1 BlockPred rep
p2 SymbolTable rep
body UsageTable
vtable Stm rep
need = BlockPred rep
p1 SymbolTable rep
body UsageTable
vtable Stm rep
need Bool -> Bool -> Bool
&& BlockPred rep
p2 SymbolTable rep
body UsageTable
vtable Stm rep
need
isConsumed :: BlockPred rep
isConsumed :: forall {k} (rep :: k). BlockPred rep
isConsumed SymbolTable rep
_ UsageTable
utable = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> UsageTable -> Bool
`UT.isConsumed` UsageTable
utable) 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 {k} (rep :: k). Stm rep -> Pat (LetDec rep)
stmPat
isOp :: BlockPred rep
isOp :: forall {k} (rep :: k). BlockPred rep
isOp SymbolTable rep
_ UsageTable
_ (Let Pat (LetDec rep)
_ StmAux (ExpDec rep)
_ Op {}) = Bool
True
isOp SymbolTable rep
_ UsageTable
_ Stm rep
_ = Bool
False
constructBody ::
SimplifiableRep rep =>
Stms (Wise rep) ->
Result ->
SimpleM rep (Body (Wise rep))
constructBody :: forall {k} (rep :: k).
SimplifiableRep rep =>
Stms (Wise rep) -> Result -> SimpleM rep (Body (Wise rep))
constructBody Stms (Wise rep)
stms Result
res =
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k1} {k2} (m :: * -> *) (somerep :: k1) (rep :: k2) a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
MonadBuilder m =>
m Result -> m (Body (Rep m))
buildBody_ forall a b. (a -> b) -> a -> b
$ do
forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms Stms (Wise rep)
stms
forall (f :: * -> *) a. Applicative f => a -> f a
pure Result
res
blockIf ::
SimplifiableRep rep =>
BlockPred (Wise rep) ->
Stms (Wise rep) ->
SimpleM rep (a, UT.UsageTable) ->
SimpleM rep (a, Stms (Wise rep), Stms (Wise rep))
blockIf :: forall {k} (rep :: k) a.
SimplifiableRep rep =>
BlockPred (Wise rep)
-> Stms (Wise rep)
-> SimpleM rep (a, UsageTable)
-> SimpleM rep (a, Stms (Wise rep), Stms (Wise rep))
blockIf BlockPred (Wise rep)
block Stms (Wise rep)
stms SimpleM rep (a, UsageTable)
m = do
RuleBook (Wise rep)
rules <- forall {k} (rep :: k) a. (Env rep -> a) -> SimpleM rep a
asksEngineEnv forall {k} (rep :: k). Env rep -> RuleBook (Wise rep)
envRules
forall {k} (rep :: k) a.
SimplifiableRep rep =>
RuleBook (Wise rep)
-> BlockPred (Wise rep)
-> Stms (Wise rep)
-> SimpleM rep (a, UsageTable)
-> SimpleM rep (a, Stms (Wise rep), Stms (Wise rep))
hoistStms RuleBook (Wise rep)
rules BlockPred (Wise rep)
block Stms (Wise rep)
stms SimpleM rep (a, UsageTable)
m
hasFree :: ASTRep rep => Names -> BlockPred rep
hasFree :: forall {k} (rep :: k). ASTRep rep => Names -> BlockPred rep
hasFree Names
ks SymbolTable rep
_ UsageTable
_ Stm rep
need = Names
ks Names -> Names -> Bool
`namesIntersect` forall a. FreeIn a => a -> Names
freeIn Stm rep
need
isNotSafe :: ASTRep rep => BlockPred rep
isNotSafe :: forall {k} (rep :: k). ASTRep rep => BlockPred rep
isNotSafe SymbolTable rep
_ UsageTable
_ = Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). IsOp (Op rep) => Exp rep -> Bool
safeExp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). Stm rep -> Exp rep
stmExp
isConsuming :: Aliased rep => BlockPred rep
isConsuming :: forall {k} (rep :: k). Aliased rep => BlockPred rep
isConsuming SymbolTable rep
_ UsageTable
_ = forall {k} {rep :: k}. Aliased rep => Exp rep -> Bool
isUpdate forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). Stm rep -> Exp rep
stmExp
where
isUpdate :: Exp rep -> Bool
isUpdate Exp rep
e = forall {k} (rep :: k). Aliased rep => Exp rep -> Names
consumedInExp Exp rep
e forall a. Eq a => a -> a -> Bool
/= forall a. Monoid a => a
mempty
isNotCheap :: ASTRep rep => BlockPred rep
isNotCheap :: forall {k} (rep :: k). ASTRep rep => BlockPred rep
isNotCheap SymbolTable rep
_ UsageTable
_ = Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). ASTRep rep => Stm rep -> Bool
cheapStm
cheapStm :: ASTRep rep => Stm rep -> Bool
cheapStm :: forall {k} (rep :: k). ASTRep rep => Stm rep -> Bool
cheapStm = forall {k} {rep :: k}. ASTRep rep => Exp rep -> Bool
cheapExp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). Stm rep -> Exp rep
stmExp
cheapExp :: ASTRep rep => Exp rep -> Bool
cheapExp :: forall {k} {rep :: k}. ASTRep rep => Exp rep -> Bool
cheapExp (BasicOp BinOp {}) = Bool
True
cheapExp (BasicOp SubExp {}) = Bool
True
cheapExp (BasicOp UnOp {}) = Bool
True
cheapExp (BasicOp CmpOp {}) = Bool
True
cheapExp (BasicOp ConvOp {}) = Bool
True
cheapExp (BasicOp Assert {}) = Bool
True
cheapExp (BasicOp Copy {}) = Bool
False
cheapExp (BasicOp Replicate {}) = Bool
False
cheapExp (BasicOp Concat {}) = Bool
False
cheapExp (BasicOp Manifest {}) = Bool
False
cheapExp DoLoop {} = Bool
False
cheapExp (Match [SubExp]
_ [Case (Body rep)]
cases Body rep
defbranch MatchDec (BranchType rep)
_) =
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall {k} (rep :: k). ASTRep rep => Stm rep -> Bool
cheapStm forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). Body rep -> Stms rep
bodyStms forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall body. Case body -> body
caseBody) [Case (Body rep)]
cases
Bool -> Bool -> Bool
&& forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall {k} (rep :: k). ASTRep rep => Stm rep -> Bool
cheapStm (forall {k} (rep :: k). Body rep -> Stms rep
bodyStms Body rep
defbranch)
cheapExp (Op Op rep
op) = forall op. IsOp op => op -> Bool
cheapOp Op rep
op
cheapExp Exp rep
_ = Bool
True
loopInvariantStm :: ASTRep rep => ST.SymbolTable rep -> Stm rep -> Bool
loopInvariantStm :: forall {k} (rep :: k).
ASTRep rep =>
SymbolTable rep -> Stm rep -> Bool
loopInvariantStm SymbolTable rep
vtable =
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> Names -> Bool
`nameIn` forall {k} (rep :: k). SymbolTable rep -> Names
ST.availableAtClosestLoop SymbolTable rep
vtable) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Names -> [VName]
namesToList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. FreeIn a => a -> Names
freeIn
matchBlocker ::
(ASTRep rep, CanBeWise (Op rep), FreeIn a) =>
a ->
MatchDec rt ->
SimpleM rep (BlockPred (Wise rep))
matchBlocker :: forall {k} (rep :: k) a rt.
(ASTRep rep, CanBeWise (Op rep), FreeIn a) =>
a -> MatchDec rt -> SimpleM rep (BlockPred (Wise rep))
matchBlocker a
cond (MatchDec [rt]
_ MatchSort
ifsort) = do
Stm (Wise rep) -> Bool
is_alloc_fun <- forall {k} (rep :: k) a. (Env rep -> a) -> SimpleM rep a
asksEngineEnv forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). HoistBlockers rep -> Stm (Wise rep) -> Bool
isAllocation forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). Env rep -> HoistBlockers rep
envHoistBlockers
BlockPred (Wise rep)
branch_blocker <- forall {k} (rep :: k) a. (Env rep -> a) -> SimpleM rep a
asksEngineEnv forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). HoistBlockers rep -> BlockPred (Wise rep)
blockHoistBranch forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). Env rep -> HoistBlockers rep
envHoistBlockers
SymbolTable (Wise rep)
vtable <- forall {k} (rep :: k). SimpleM rep (SymbolTable (Wise rep))
askVtable
let
cond_loop_invariant :: Bool
cond_loop_invariant =
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> Names -> Bool
`nameIn` forall {k} (rep :: k). SymbolTable rep -> Names
ST.availableAtClosestLoop SymbolTable (Wise rep)
vtable) forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn a
cond
desirableToHoist :: UsageTable -> Stm (Wise rep) -> Bool
desirableToHoist UsageTable
usage Stm (Wise rep)
stm =
Stm (Wise rep) -> Bool
is_alloc_fun Stm (Wise rep)
stm
Bool -> Bool -> Bool
|| ( forall {k} (rep :: k). SymbolTable rep -> Int
ST.loopDepth SymbolTable (Wise rep)
vtable forall a. Ord a => a -> a -> Bool
> Int
0
Bool -> Bool -> Bool
&& Bool
cond_loop_invariant
Bool -> Bool -> Bool
&& MatchSort
ifsort forall a. Eq a => a -> a -> Bool
/= MatchSort
MatchFallback
Bool -> Bool -> Bool
&& forall {k} (rep :: k).
ASTRep rep =>
SymbolTable rep -> Stm rep -> Bool
loopInvariantStm SymbolTable (Wise rep)
vtable Stm (Wise rep)
stm
Bool -> Bool -> Bool
&& forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall shape u. TypeBase shape u -> Bool
primType (forall dec. Typed dec => Pat dec -> [Type]
patTypes (forall {k} (rep :: k). Stm rep -> Pat (LetDec rep)
stmPat Stm (Wise rep)
stm))
)
Bool -> Bool -> Bool
|| ( MatchSort
ifsort forall a. Eq a => a -> a -> Bool
/= MatchSort
MatchFallback
Bool -> Bool -> Bool
&& forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> UsageTable -> Bool
`UT.isSize` UsageTable
usage) (forall dec. Pat dec -> [VName]
patNames (forall {k} (rep :: k). Stm rep -> Pat (LetDec rep)
stmPat Stm (Wise rep)
stm))
Bool -> Bool -> Bool
&& forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall shape u. TypeBase shape u -> Bool
primType (forall dec. Typed dec => Pat dec -> [Type]
patTypes (forall {k} (rep :: k). Stm rep -> Pat (LetDec rep)
stmPat Stm (Wise rep)
stm))
)
notDesirableToHoist :: BlockPred (Wise rep)
notDesirableToHoist SymbolTable (Wise rep)
_ UsageTable
usage Stm (Wise rep)
stm = Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ UsageTable -> Stm (Wise rep) -> Bool
desirableToHoist UsageTable
usage Stm (Wise rep)
stm
isNotHoistableBnd :: BlockPred (Wise rep)
isNotHoistableBnd SymbolTable (Wise rep)
_ UsageTable
_ (Let Pat (LetDec (Wise rep))
_ StmAux (ExpDec (Wise rep))
_ (BasicOp ArrayLit {})) = Bool
False
isNotHoistableBnd SymbolTable (Wise rep)
_ UsageTable
_ (Let Pat (LetDec (Wise rep))
_ StmAux (ExpDec (Wise rep))
_ (BasicOp SubExp {})) = Bool
False
isNotHoistableBnd SymbolTable (Wise rep)
_ UsageTable
_ (Let Pat (LetDec (Wise rep))
_ StmAux (ExpDec (Wise rep))
_ (BasicOp Reshape {})) = Bool
False
isNotHoistableBnd SymbolTable (Wise rep)
_ UsageTable
_ (Let Pat (LetDec (Wise rep))
_ StmAux (ExpDec (Wise rep))
_ (BasicOp Rearrange {})) = Bool
False
isNotHoistableBnd SymbolTable (Wise rep)
_ UsageTable
_ (Let Pat (LetDec (Wise rep))
_ StmAux (ExpDec (Wise rep))
_ (BasicOp Rotate {})) = Bool
False
isNotHoistableBnd SymbolTable (Wise rep)
_ UsageTable
_ (Let Pat (LetDec (Wise rep))
_ StmAux (ExpDec (Wise rep))
_ (BasicOp (Index VName
_ Slice SubExp
slice))) =
forall (t :: * -> *) a. Foldable t => t a -> Bool
null forall a b. (a -> b) -> a -> b
$ forall d. Slice d -> [d]
sliceDims Slice SubExp
slice
isNotHoistableBnd SymbolTable (Wise rep)
_ UsageTable
_ Stm (Wise rep)
stm
| Stm (Wise rep) -> Bool
is_alloc_fun Stm (Wise rep)
stm = Bool
False
isNotHoistableBnd SymbolTable (Wise rep)
_ UsageTable
_ Stm (Wise rep)
_ =
MatchSort
ifsort forall a. Eq a => a -> a -> Bool
/= MatchSort
MatchEquiv
block :: BlockPred (Wise rep)
block =
BlockPred (Wise rep)
branch_blocker
forall {k} (rep :: k).
BlockPred rep -> BlockPred rep -> BlockPred rep
`orIf` ( (forall {k} (rep :: k). ASTRep rep => BlockPred rep
isNotSafe forall {k} (rep :: k).
BlockPred rep -> BlockPred rep -> BlockPred rep
`orIf` forall {k} (rep :: k). ASTRep rep => BlockPred rep
isNotCheap forall {k} (rep :: k).
BlockPred rep -> BlockPred rep -> BlockPred rep
`orIf` BlockPred (Wise rep)
isNotHoistableBnd)
forall {k} (rep :: k).
BlockPred rep -> BlockPred rep -> BlockPred rep
`andAlso` BlockPred (Wise rep)
notDesirableToHoist
)
forall {k} (rep :: k).
BlockPred rep -> BlockPred rep -> BlockPred rep
`orIf` forall {k} (rep :: k). Aliased rep => BlockPred rep
isConsuming
forall (f :: * -> *) a. Applicative f => a -> f a
pure BlockPred (Wise rep)
block
simplifyBody ::
SimplifiableRep rep =>
BlockPred (Wise rep) ->
UT.UsageTable ->
[UT.Usages] ->
Body (Wise rep) ->
SimpleM rep (Stms (Wise rep), Body (Wise rep))
simplifyBody :: forall {k} (rep :: k).
SimplifiableRep rep =>
BlockPred (Wise rep)
-> UsageTable
-> [Usages]
-> Body (Wise rep)
-> SimpleM rep (Stms (Wise rep), Body (Wise rep))
simplifyBody BlockPred (Wise rep)
blocker UsageTable
usage [Usages]
res_usages (Body BodyDec (Wise rep)
_ Stms (Wise rep)
stms Result
res) = do
(Result
res', Stms (Wise rep)
stms', Stms (Wise rep)
hoisted) <-
forall {k} (rep :: k) a.
SimplifiableRep rep =>
BlockPred (Wise rep)
-> Stms (Wise rep)
-> SimpleM rep (a, UsageTable)
-> SimpleM rep (a, Stms (Wise rep), Stms (Wise rep))
blockIf BlockPred (Wise rep)
blocker Stms (Wise rep)
stms forall a b. (a -> b) -> a -> b
$ do
(Result
res', UsageTable
res_usage) <- forall {k} (rep :: k).
SimplifiableRep rep =>
[Usages] -> Result -> SimpleM rep (Result, UsageTable)
simplifyResult [Usages]
res_usages Result
res
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Result
res', UsageTable
res_usage forall a. Semigroup a => a -> a -> a
<> UsageTable
usage)
Body (Wise rep)
body' <- forall {k} (rep :: k).
SimplifiableRep rep =>
Stms (Wise rep) -> Result -> SimpleM rep (Body (Wise rep))
constructBody Stms (Wise rep)
stms' Result
res'
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms (Wise rep)
hoisted, Body (Wise rep)
body')
simplifyBodyNoHoisting ::
SimplifiableRep rep =>
UT.UsageTable ->
[UT.Usages] ->
Body (Wise rep) ->
SimpleM rep (Body (Wise rep))
simplifyBodyNoHoisting :: forall {k} (rep :: k).
SimplifiableRep rep =>
UsageTable
-> [Usages] -> Body (Wise rep) -> SimpleM rep (Body (Wise rep))
simplifyBodyNoHoisting UsageTable
usage [Usages]
res_usages Body (Wise rep)
body =
forall a b. (a, b) -> b
snd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (rep :: k).
SimplifiableRep rep =>
BlockPred (Wise rep)
-> UsageTable
-> [Usages]
-> Body (Wise rep)
-> SimpleM rep (Stms (Wise rep), Body (Wise rep))
simplifyBody (forall {k} (rep :: k). Bool -> BlockPred rep
isFalse Bool
False) UsageTable
usage [Usages]
res_usages Body (Wise rep)
body
usageFromDiet :: Diet -> UT.Usages
usageFromDiet :: Diet -> Usages
usageFromDiet Diet
Consume = Usages
UT.consumedU
usageFromDiet Diet
_ = forall a. Monoid a => a
mempty
simplifyResult ::
SimplifiableRep rep => [UT.Usages] -> Result -> SimpleM rep (Result, UT.UsageTable)
simplifyResult :: forall {k} (rep :: k).
SimplifiableRep rep =>
[Usages] -> Result -> SimpleM rep (Result, UsageTable)
simplifyResult [Usages]
usages Result
res = do
Result
res' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify Result
res
SymbolTable (Wise rep)
vtable <- forall {k} (rep :: k). SimpleM rep (SymbolTable (Wise rep))
askVtable
let more_usages :: UsageTable
more_usages = forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ do
(Usages
u, Var VName
v) <- forall a b. [a] -> [b] -> [(a, b)]
zip [Usages]
usages forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp Result
res
let als_usages :: [UsageTable]
als_usages =
forall a b. (a -> b) -> [a] -> [b]
map
(VName -> Usages -> UsageTable
`UT.usage` (Usages
u Usages -> Usages -> Usages
`UT.withoutU` Usages
UT.presentU))
(Names -> [VName]
namesToList (forall {k} (rep :: k). VName -> SymbolTable rep -> Names
ST.lookupAliases VName
v SymbolTable (Wise rep)
vtable))
VName -> Usages -> UsageTable
UT.usage VName
v Usages
u forall a. a -> [a] -> [a]
: [UsageTable]
als_usages
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( Result
res',
Names -> UsageTable
UT.usages (forall a. FreeIn a => a -> Names
freeIn Result
res')
forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap VName -> UsageTable
UT.inResultUsage (Names -> [VName]
namesToList (forall a. FreeIn a => a -> Names
freeIn Result
res'))
forall a. Semigroup a => a -> a -> a
<> UsageTable
more_usages
)
isDoLoopResult :: Result -> UT.UsageTable
isDoLoopResult :: Result -> UsageTable
isDoLoopResult = 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 SubExpRes -> UsageTable
checkForVar
where
checkForVar :: SubExpRes -> UsageTable
checkForVar (SubExpRes Certs
_ (Var VName
ident)) = VName -> UsageTable
UT.inResultUsage VName
ident
checkForVar SubExpRes
_ = forall a. Monoid a => a
mempty
simplifyStms ::
SimplifiableRep rep =>
Stms (Wise rep) ->
SimpleM rep (Stms (Wise rep))
simplifyStms :: forall {k} (rep :: k).
SimplifiableRep rep =>
Stms (Wise rep) -> SimpleM rep (Stms (Wise rep))
simplifyStms Stms (Wise rep)
stms = do
forall {k} (rep :: k).
SimplifiableRep rep =>
UsageTable -> Stms (Wise rep) -> SimpleM rep (Stms (Wise rep))
simplifyStmsWithUsage UsageTable
all_used Stms (Wise rep)
stms
where
all_used :: UsageTable
all_used =
Names -> UsageTable
UT.usages ([VName] -> Names
namesFromList (forall k a. Map k a -> [k]
M.keys (forall {k} (rep :: k) a. Scoped rep a => a -> Scope rep
scopeOf Stms (Wise rep)
stms)))
simplifyStmsWithUsage ::
SimplifiableRep rep =>
UT.UsageTable ->
Stms (Wise rep) ->
SimpleM rep (Stms (Wise rep))
simplifyStmsWithUsage :: forall {k} (rep :: k).
SimplifiableRep rep =>
UsageTable -> Stms (Wise rep) -> SimpleM rep (Stms (Wise rep))
simplifyStmsWithUsage UsageTable
usage Stms (Wise rep)
stms = do
((), Stms (Wise rep)
stms', Stms (Wise rep)
_) <- forall {k} (rep :: k) a.
SimplifiableRep rep =>
BlockPred (Wise rep)
-> Stms (Wise rep)
-> SimpleM rep (a, UsageTable)
-> SimpleM rep (a, Stms (Wise rep), Stms (Wise rep))
blockIf (forall {k} (rep :: k). Bool -> BlockPred rep
isFalse Bool
False) Stms (Wise rep)
stms forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure ((), UsageTable
usage)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Stms (Wise rep)
stms'
simplifyOp :: Op (Wise rep) -> SimpleM rep (Op (Wise rep), Stms (Wise rep))
simplifyOp :: forall {k} (rep :: k).
Op (Wise rep) -> SimpleM rep (Op (Wise rep), Stms (Wise rep))
simplifyOp Op (Wise rep)
op = do
OpWithWisdom (Op rep)
-> SimpleM rep (OpWithWisdom (Op rep), Stms (Wise rep))
f <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
SimpleOps rep -> SimplifyOp rep (Op (Wise rep))
simplifyOpS forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst
OpWithWisdom (Op rep)
-> SimpleM rep (OpWithWisdom (Op rep), Stms (Wise rep))
f Op (Wise rep)
op
simplifyExp ::
SimplifiableRep rep =>
UT.UsageTable ->
Pat (LetDec (Wise rep)) ->
Exp (Wise rep) ->
SimpleM rep (Exp (Wise rep), Stms (Wise rep))
simplifyExp :: forall {k} (rep :: k).
SimplifiableRep rep =>
UsageTable
-> Pat (LetDec (Wise rep))
-> Exp (Wise rep)
-> SimpleM rep (Exp (Wise rep), Stms (Wise rep))
simplifyExp UsageTable
usage (Pat [PatElem (LetDec (Wise rep))]
pes) (Match [SubExp]
ses [Case (Body (Wise rep))]
cases Body (Wise rep)
defbody ifdec :: MatchDec (BranchType (Wise rep))
ifdec@(MatchDec [BranchType (Wise rep)]
ts MatchSort
ifsort)) = do
let pes_usages :: [Usages]
pes_usages = forall a b. (a -> b) -> [a] -> [b]
map (forall a. a -> Maybe a -> a
fromMaybe forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> UsageTable -> Maybe Usages
`UT.lookup` UsageTable
usage) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. PatElem dec -> VName
patElemName) [PatElem (LetDec (Wise rep))]
pes
[SubExp]
ses' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify [SubExp]
ses
[BranchType rep]
ts' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify [BranchType (Wise rep)]
ts
let pats :: [[Maybe PrimValue]]
pats = forall a b. (a -> b) -> [a] -> [b]
map forall body. Case body -> [Maybe PrimValue]
casePat [Case (Body (Wise rep))]
cases
BlockPred (Wise rep)
block <- forall {k} (rep :: k) a rt.
(ASTRep rep, CanBeWise (Op rep), FreeIn a) =>
a -> MatchDec rt -> SimpleM rep (BlockPred (Wise rep))
matchBlocker [SubExp]
ses MatchDec (BranchType (Wise rep))
ifdec
([Stms (Wise rep)]
cases_hoisted, [Case (Body (Wise rep))]
cases') <-
forall a b. [(a, b)] -> ([a], [b])
unzip forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (BlockPred (Wise rep)
-> [SubExp]
-> [Usages]
-> [[Maybe PrimValue]]
-> Case (Body (Wise rep))
-> SimpleM rep (Stms (Wise rep), Case (Body (Wise rep)))
simplifyCase BlockPred (Wise rep)
block [SubExp]
ses' [Usages]
pes_usages) (forall a. [a] -> [[a]]
inits [[Maybe PrimValue]]
pats) [Case (Body (Wise rep))]
cases
(Stms (Wise rep)
defbody_hoisted, Body (Wise rep)
defbody') <-
forall {k} (rep :: k) a.
SimplifiableRep rep =>
[SubExp]
-> [[Maybe PrimValue]]
-> [Maybe PrimValue]
-> SimpleM rep (Stms (Wise rep), a)
-> SimpleM rep (Stms (Wise rep), a)
protectCaseHoisted [SubExp]
ses' [[Maybe PrimValue]]
pats [] forall a b. (a -> b) -> a -> b
$
forall {k} (rep :: k).
SimplifiableRep rep =>
BlockPred (Wise rep)
-> UsageTable
-> [Usages]
-> Body (Wise rep)
-> SimpleM rep (Stms (Wise rep), Body (Wise rep))
simplifyBody BlockPred (Wise rep)
block UsageTable
usage [Usages]
pes_usages Body (Wise rep)
defbody
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall {k} (rep :: k).
[SubExp]
-> [Case (Body rep)]
-> Body rep
-> MatchDec (BranchType rep)
-> Exp rep
Match [SubExp]
ses' [Case (Body (Wise rep))]
cases' Body (Wise rep)
defbody' forall a b. (a -> b) -> a -> b
$ forall rt. [rt] -> MatchSort -> MatchDec rt
MatchDec [BranchType rep]
ts' MatchSort
ifsort,
forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ Stms (Wise rep)
defbody_hoisted forall a. a -> [a] -> [a]
: [Stms (Wise rep)]
cases_hoisted
)
where
simplifyCase :: BlockPred (Wise rep)
-> [SubExp]
-> [Usages]
-> [[Maybe PrimValue]]
-> Case (Body (Wise rep))
-> SimpleM rep (Stms (Wise rep), Case (Body (Wise rep)))
simplifyCase BlockPred (Wise rep)
block [SubExp]
ses' [Usages]
pes_usages [[Maybe PrimValue]]
prior (Case [Maybe PrimValue]
vs Body (Wise rep)
body) = do
(Stms (Wise rep)
hoisted, Body (Wise rep)
body') <-
forall {k} (rep :: k) a.
SimplifiableRep rep =>
[SubExp]
-> [[Maybe PrimValue]]
-> [Maybe PrimValue]
-> SimpleM rep (Stms (Wise rep), a)
-> SimpleM rep (Stms (Wise rep), a)
protectCaseHoisted [SubExp]
ses' [[Maybe PrimValue]]
prior [Maybe PrimValue]
vs forall a b. (a -> b) -> a -> b
$
forall {k} (rep :: k).
SimplifiableRep rep =>
BlockPred (Wise rep)
-> UsageTable
-> [Usages]
-> Body (Wise rep)
-> SimpleM rep (Stms (Wise rep), Body (Wise rep))
simplifyBody BlockPred (Wise rep)
block UsageTable
usage [Usages]
pes_usages Body (Wise rep)
body
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms (Wise rep)
hoisted, forall body. [Maybe PrimValue] -> body -> Case body
Case [Maybe PrimValue]
vs Body (Wise rep)
body')
simplifyExp UsageTable
_ Pat (LetDec (Wise rep))
_ (DoLoop [(FParam (Wise rep), SubExp)]
merge LoopForm (Wise rep)
form Body (Wise rep)
loopbody) = do
let ([Param (FParamInfo rep)]
params, [SubExp]
args) = forall a b. [(a, b)] -> ([a], [b])
unzip [(FParam (Wise rep), SubExp)]
merge
[Param (FParamInfo rep)]
params' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify) [Param (FParamInfo rep)]
params
[SubExp]
args' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify [SubExp]
args
let merge' :: [(Param (FParamInfo rep), SubExp)]
merge' = forall a b. [a] -> [b] -> [(a, b)]
zip [Param (FParamInfo rep)]
params' [SubExp]
args'
(LoopForm (Wise rep)
form', Names
boundnames, SimpleM rep (Result, Stms (Wise rep), Stms (Wise rep))
-> SimpleM rep (Result, Stms (Wise rep), Stms (Wise rep))
wrapbody) <- case LoopForm (Wise rep)
form of
ForLoop VName
loopvar IntType
it SubExp
boundexp [(LParam (Wise rep), VName)]
loopvars -> do
SubExp
boundexp' <- forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify SubExp
boundexp
let ([Param (LParamInfo rep)]
loop_params, [VName]
loop_arrs) = forall a b. [(a, b)] -> ([a], [b])
unzip [(LParam (Wise rep), VName)]
loopvars
[Param (LParamInfo rep)]
loop_params' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify) [Param (LParamInfo rep)]
loop_params
[VName]
loop_arrs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify [VName]
loop_arrs
let form' :: LoopForm (Wise rep)
form' = forall {k} (rep :: k).
VName -> IntType -> SubExp -> [(LParam rep, VName)] -> LoopForm rep
ForLoop VName
loopvar IntType
it SubExp
boundexp' (forall a b. [a] -> [b] -> [(a, b)]
zip [Param (LParamInfo rep)]
loop_params' [VName]
loop_arrs')
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( LoopForm (Wise rep)
form',
[VName] -> Names
namesFromList (VName
loopvar forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
paramName [Param (LParamInfo rep)]
loop_params') forall a. Semigroup a => a -> a -> a
<> Names
fparamnames,
forall {k} (rep :: k) a.
SimplifiableRep rep =>
VName -> IntType -> SubExp -> SimpleM rep a -> SimpleM rep a
bindLoopVar VName
loopvar IntType
it SubExp
boundexp'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k) a b.
SimplifiableRep rep =>
[(FParam (Wise rep), SubExp)]
-> LoopForm (Wise rep)
-> SimpleM rep (a, b, Stms (Wise rep))
-> SimpleM rep (a, b, Stms (Wise rep))
protectLoopHoisted [(Param (FParamInfo rep), SubExp)]
merge' LoopForm (Wise rep)
form'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k) a.
SimplifiableRep rep =>
[LParam (Wise rep)] -> SimpleM rep a -> SimpleM rep a
bindArrayLParams [Param (LParamInfo rep)]
loop_params'
)
WhileLoop VName
cond -> do
VName
cond' <- forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify VName
cond
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall {k} (rep :: k). VName -> LoopForm rep
WhileLoop VName
cond',
Names
fparamnames,
forall {k} (rep :: k) a b.
SimplifiableRep rep =>
[(FParam (Wise rep), SubExp)]
-> LoopForm (Wise rep)
-> SimpleM rep (a, b, Stms (Wise rep))
-> SimpleM rep (a, b, Stms (Wise rep))
protectLoopHoisted [(Param (FParamInfo rep), SubExp)]
merge' (forall {k} (rep :: k). VName -> LoopForm rep
WhileLoop VName
cond')
)
BlockPred (Wise rep)
seq_blocker <- forall {k} (rep :: k) a. (Env rep -> a) -> SimpleM rep a
asksEngineEnv forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). HoistBlockers rep -> BlockPred (Wise rep)
blockHoistSeq forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). Env rep -> HoistBlockers rep
envHoistBlockers
(Result
loopres, Stms (Wise rep)
loopstms, Stms (Wise rep)
hoisted) <-
forall {k} (rep :: k) a. SimpleM rep a -> SimpleM rep a
enterLoop forall b c a. (b -> c) -> (a -> b) -> a -> c
. SimpleM rep (Result, Stms (Wise rep), Stms (Wise rep))
-> SimpleM rep (Result, Stms (Wise rep), Stms (Wise rep))
consumeMerge
forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k) a.
SimplifiableRep rep =>
[(FParam (Wise rep), SubExp, SubExpRes)]
-> SimpleM rep a -> SimpleM rep a
bindMerge (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall {a} {b} {c}. (a, b) -> c -> (a, b, c)
withRes [(Param (FParamInfo rep), SubExp)]
merge' (forall {k} (rep :: k). Body rep -> Result
bodyResult Body (Wise rep)
loopbody)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. SimpleM rep (Result, Stms (Wise rep), Stms (Wise rep))
-> SimpleM rep (Result, Stms (Wise rep), Stms (Wise rep))
wrapbody
forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k) a.
SimplifiableRep rep =>
BlockPred (Wise rep)
-> Stms (Wise rep)
-> SimpleM rep (a, UsageTable)
-> SimpleM rep (a, Stms (Wise rep), Stms (Wise rep))
blockIf
( forall {k} (rep :: k). ASTRep rep => Names -> BlockPred rep
hasFree Names
boundnames
forall {k} (rep :: k).
BlockPred rep -> BlockPred rep -> BlockPred rep
`orIf` forall {k} (rep :: k). BlockPred rep
isConsumed
forall {k} (rep :: k).
BlockPred rep -> BlockPred rep -> BlockPred rep
`orIf` BlockPred (Wise rep)
seq_blocker
forall {k} (rep :: k).
BlockPred rep -> BlockPred rep -> BlockPred rep
`orIf` forall {k} (rep :: k). ASTRep rep => BlockPred rep
notWorthHoisting
)
(forall {k} (rep :: k). Body rep -> Stms rep
bodyStms Body (Wise rep)
loopbody)
forall a b. (a -> b) -> a -> b
$ do
let params_usages :: [Usages]
params_usages =
forall a b. (a -> b) -> [a] -> [b]
map
(\Param (FParamInfo rep)
p -> if forall shape. TypeBase shape Uniqueness -> Bool
unique (forall dec. DeclTyped dec => Param dec -> DeclType
paramDeclType Param (FParamInfo rep)
p) then Usages
UT.consumedU else forall a. Monoid a => a
mempty)
[Param (FParamInfo rep)]
params'
(Result
res, UsageTable
uses) <- forall {k} (rep :: k).
SimplifiableRep rep =>
[Usages] -> Result -> SimpleM rep (Result, UsageTable)
simplifyResult [Usages]
params_usages forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Body rep -> Result
bodyResult Body (Wise rep)
loopbody
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Result
res, UsageTable
uses forall a. Semigroup a => a -> a -> a
<> Result -> UsageTable
isDoLoopResult Result
res)
Body (Wise rep)
loopbody' <- forall {k} (rep :: k).
SimplifiableRep rep =>
Stms (Wise rep) -> Result -> SimpleM rep (Body (Wise rep))
constructBody Stms (Wise rep)
loopstms Result
loopres
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall {k} (rep :: k).
[(FParam rep, SubExp)] -> LoopForm rep -> Body rep -> Exp rep
DoLoop [(Param (FParamInfo rep), SubExp)]
merge' LoopForm (Wise rep)
form' Body (Wise rep)
loopbody', Stms (Wise rep)
hoisted)
where
fparamnames :: Names
fparamnames =
[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 (Wise rep), SubExp)]
merge)
consumeMerge :: SimpleM rep (Result, Stms (Wise rep), Stms (Wise rep))
-> SimpleM rep (Result, Stms (Wise rep), Stms (Wise rep))
consumeMerge =
forall {k} (rep :: k) a.
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep a -> SimpleM rep a
localVtable 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 a b c. (a -> b -> c) -> b -> a -> c
flip forall {k} (rep :: k). VName -> SymbolTable rep -> SymbolTable rep
ST.consume)) forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList Names
consumed_by_merge
consumed_by_merge :: Names
consumed_by_merge =
forall a. FreeIn a => a -> Names
freeIn forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
filter (forall shape. TypeBase shape Uniqueness -> Bool
unique forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. DeclTyped dec => Param dec -> DeclType
paramDeclType forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(FParam (Wise rep), SubExp)]
merge
withRes :: (a, b) -> c -> (a, b, c)
withRes (a
p, b
x) c
y = (a
p, b
x, c
y)
simplifyExp UsageTable
_ Pat (LetDec (Wise rep))
_ (Op Op (Wise rep)
op) = do
(OpWithWisdom (Op rep)
op', Stms (Wise rep)
stms) <- forall {k} (rep :: k).
Op (Wise rep) -> SimpleM rep (Op (Wise rep), Stms (Wise rep))
simplifyOp Op (Wise rep)
op
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall {k} (rep :: k). Op rep -> Exp rep
Op OpWithWisdom (Op rep)
op', Stms (Wise rep)
stms)
simplifyExp UsageTable
usage Pat (LetDec (Wise rep))
_ (WithAcc [WithAccInput (Wise rep)]
inputs Lambda (Wise rep)
lam) = do
([WithAccInput (Wise rep)]
inputs', [Stms (Wise rep)]
inputs_stms) <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. [(a, b)] -> ([a], [b])
unzip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [WithAccInput (Wise rep)]
inputs forall a b. (a -> b) -> a -> b
$ \(Shape
shape, [VName]
arrs, Maybe (Lambda (Wise rep), [SubExp])
op) -> do
(Maybe (Lambda (Wise rep), [SubExp])
op', Stms (Wise rep)
op_stms) <- case Maybe (Lambda (Wise rep), [SubExp])
op of
Maybe (Lambda (Wise rep), [SubExp])
Nothing ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Maybe a
Nothing, forall a. Monoid a => a
mempty)
Just (Lambda (Wise rep)
op_lam, [SubExp]
nes) -> do
(Lambda (Wise rep)
op_lam', Stms (Wise rep)
op_lam_stms) <- forall {k} (rep :: k).
SimplifiableRep rep =>
SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
blockMigrated (forall {k} (rep :: k).
SimplifiableRep rep =>
Lambda (Wise rep)
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
simplifyLambda Lambda (Wise rep)
op_lam)
[SubExp]
nes' <- forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify [SubExp]
nes
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. a -> Maybe a
Just (Lambda (Wise rep)
op_lam', [SubExp]
nes'), Stms (Wise rep)
op_lam_stms)
(,Stms (Wise rep)
op_stms) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((,,Maybe (Lambda (Wise rep), [SubExp])
op') forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify Shape
shape forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify [VName]
arrs)
let noteAcc :: SymbolTable (Wise rep) -> SymbolTable (Wise rep)
noteAcc = forall {k} (rep :: k).
[(VName, WithAccInput rep)] -> SymbolTable rep -> SymbolTable rep
ST.noteAccTokens (forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
paramName (forall {k} (rep :: k). Lambda rep -> [LParam rep]
lambdaParams Lambda (Wise rep)
lam)) [WithAccInput (Wise rep)]
inputs')
(Lambda (Wise rep)
lam', Stms (Wise rep)
lam_stms) <- forall {k} (rep :: k).
SimplifiableRep rep =>
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> BlockPred (Wise rep)
-> UsageTable
-> Lambda (Wise rep)
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
simplifyLambdaWith SymbolTable (Wise rep) -> SymbolTable (Wise rep)
noteAcc (forall {k} (rep :: k). Bool -> BlockPred rep
isFalse Bool
True) UsageTable
usage Lambda (Wise rep)
lam
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall {k} (rep :: k). [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc [WithAccInput (Wise rep)]
inputs' Lambda (Wise rep)
lam', forall a. Monoid a => [a] -> a
mconcat [Stms (Wise rep)]
inputs_stms forall a. Semigroup a => a -> a -> a
<> Stms (Wise rep)
lam_stms)
simplifyExp UsageTable
_ Pat (LetDec (Wise rep))
_ Exp (Wise rep)
e = do
Exp (Wise rep)
e' <- forall {k} (rep :: k).
SimplifiableRep rep =>
Exp (Wise rep) -> SimpleM rep (Exp (Wise rep))
simplifyExpBase Exp (Wise rep)
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Wise rep)
e', forall a. Monoid a => a
mempty)
blockMigrated ::
SimplifiableRep rep =>
SimpleM rep (Lambda (Wise rep), Stms (Wise rep)) ->
SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
blockMigrated :: forall {k} (rep :: k).
SimplifiableRep rep =>
SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
blockMigrated = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall {k} {rep :: k} {a}.
(ASTRep rep, Simplifiable (LetDec rep),
Simplifiable (FParamInfo rep), Simplifiable (LParamInfo rep),
Simplifiable (RetType rep), Simplifiable (BranchType rep),
TraverseOpStms (Wise rep), CanBeWise (Op rep),
IndexOp (OpWithWisdom (Op rep)), BuilderOps (Wise rep)) =>
(a, Env rep) -> (a, Env rep)
withMigrationBlocker
where
withMigrationBlocker :: (a, Env rep) -> (a, Env rep)
withMigrationBlocker (a
ops, Env rep
env) =
let blockers :: HoistBlockers rep
blockers = forall {k} (rep :: k). Env rep -> HoistBlockers rep
envHoistBlockers Env rep
env
par_blocker :: BlockPred (Wise rep)
par_blocker = forall {k} (rep :: k). HoistBlockers rep -> BlockPred (Wise rep)
blockHoistPar HoistBlockers rep
blockers
blocker :: BlockPred (Wise rep)
blocker = BlockPred (Wise rep)
par_blocker forall {k} (rep :: k).
BlockPred rep -> BlockPred rep -> BlockPred rep
`orIf` forall {k} (rep :: k). SimplifiableRep rep => BlockPred (Wise rep)
isDeviceMigrated
blockers' :: HoistBlockers rep
blockers' = HoistBlockers rep
blockers {blockHoistPar :: BlockPred (Wise rep)
blockHoistPar = BlockPred (Wise rep)
blocker}
env' :: Env rep
env' = Env rep
env {envHoistBlockers :: HoistBlockers rep
envHoistBlockers = HoistBlockers rep
blockers'}
in (a
ops, Env rep
env')
isDeviceMigrated :: SimplifiableRep rep => BlockPred (Wise rep)
isDeviceMigrated :: forall {k} (rep :: k). SimplifiableRep rep => BlockPred (Wise rep)
isDeviceMigrated SymbolTable (Wise rep)
vtable UsageTable
_ Stm (Wise rep)
stm
| BasicOp (Index VName
arr Slice SubExp
slice) <- forall {k} (rep :: k). Stm rep -> Exp rep
stmExp Stm (Wise rep)
stm,
[DimFix SubExp
idx] <- forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slice,
SubExp
idx forall a. Eq a => a -> a -> Bool
== IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0,
Just Type
arr_t <- forall {k} (rep :: k).
ASTRep rep =>
VName -> SymbolTable rep -> Maybe Type
ST.lookupType VName
arr SymbolTable (Wise rep)
vtable,
[SubExp
size] <- forall u. TypeBase Shape u -> [SubExp]
arrayDims Type
arr_t,
SubExp
size forall a. Eq a => a -> a -> Bool
== IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1 =
Bool
True
| Bool
otherwise =
Bool
False
simplifyExpBase :: SimplifiableRep rep => Exp (Wise rep) -> SimpleM rep (Exp (Wise rep))
simplifyExpBase :: forall {k} (rep :: k).
SimplifiableRep rep =>
Exp (Wise rep) -> SimpleM rep (Exp (Wise rep))
simplifyExpBase (BasicOp (BinOp BinOp
op SubExp
x SubExp
y))
| BinOp -> Bool
commutativeBinOp BinOp
op = do
SubExp
x' <- forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify SubExp
x
SubExp
y' <- forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify SubExp
y
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp BinOp
op (forall a. Ord a => a -> a -> a
min SubExp
x' SubExp
y') (forall a. Ord a => a -> a -> a
max SubExp
x' SubExp
y')
simplifyExpBase Exp (Wise rep)
e = forall {k1} {k2} (m :: * -> *) (frep :: k1) (trep :: k2).
Monad m =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper (Wise rep) (Wise rep) (SimpleM rep)
hoist Exp (Wise rep)
e
where
hoist :: Mapper (Wise rep) (Wise rep) (SimpleM rep)
hoist =
forall {k} (m :: * -> *) (rep :: k). Monad m => Mapper rep rep m
identityMapper
{ mapOnSubExp :: SubExp -> SimpleM rep SubExp
mapOnSubExp = forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify,
mapOnVName :: VName -> SimpleM rep VName
mapOnVName = forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify,
mapOnRetType :: RetType (Wise rep) -> SimpleM rep (RetType (Wise rep))
mapOnRetType = forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify,
mapOnBranchType :: BranchType (Wise rep) -> SimpleM rep (BranchType (Wise rep))
mapOnBranchType = forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify
}
type SimplifiableRep rep =
( ASTRep rep,
Simplifiable (LetDec rep),
Simplifiable (FParamInfo rep),
Simplifiable (LParamInfo rep),
Simplifiable (RetType rep),
Simplifiable (BranchType rep),
TraverseOpStms (Wise rep),
CanBeWise (Op rep),
ST.IndexOp (OpWithWisdom (Op rep)),
BuilderOps (Wise rep),
IsOp (Op rep)
)
class Simplifiable e where
simplify :: SimplifiableRep rep => e -> SimpleM rep e
instance (Simplifiable a, Simplifiable b) => Simplifiable (a, b) where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
(a, b) -> SimpleM rep (a, b)
simplify (a
x, b
y) = (,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify a
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify b
y
instance
(Simplifiable a, Simplifiable b, Simplifiable c) =>
Simplifiable (a, b, c)
where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
(a, b, c) -> SimpleM rep (a, b, c)
simplify (a
x, b
y, c
z) = (,,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify a
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify b
y forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify c
z
instance Simplifiable Int where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
Int -> SimpleM rep Int
simplify = forall (f :: * -> *) a. Applicative f => a -> f a
pure
instance Simplifiable a => Simplifiable (Maybe a) where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
Maybe a -> SimpleM rep (Maybe a)
simplify Maybe a
Nothing = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
simplify (Just a
x) = forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify a
x
instance Simplifiable a => Simplifiable [a] where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
[a] -> SimpleM rep [a]
simplify = forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify
instance Simplifiable SubExp where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
SubExp -> SimpleM rep SubExp
simplify (Var VName
name) = do
Maybe (SubExp, Certs)
stm <- forall {k} (rep :: k).
VName -> SymbolTable rep -> Maybe (SubExp, Certs)
ST.lookupSubExp VName
name forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (rep :: k). SimpleM rep (SymbolTable (Wise rep))
askVtable
case Maybe (SubExp, Certs)
stm of
Just (Constant PrimValue
v, Certs
cs) -> do
forall {k} (rep :: k). SimpleM rep ()
changed
forall {k} (rep :: k). Certs -> SimpleM rep ()
usedCerts Certs
cs
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant PrimValue
v
Just (Var VName
id', Certs
cs) -> do
forall {k} (rep :: k). SimpleM rep ()
changed
forall {k} (rep :: k). Certs -> SimpleM rep ()
usedCerts Certs
cs
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
id'
Maybe (SubExp, Certs)
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
name
simplify (Constant PrimValue
v) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant PrimValue
v
instance Simplifiable SubExpRes where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
SubExpRes -> SimpleM rep SubExpRes
simplify (SubExpRes Certs
cs SubExp
se) = do
Certs
cs' <- forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify Certs
cs
(SubExp
se', Certs
se_cs) <- forall {k} (rep :: k) a. SimpleM rep a -> SimpleM rep (a, Certs)
collectCerts forall a b. (a -> b) -> a -> b
$ forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify SubExp
se
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Certs -> SubExp -> SubExpRes
SubExpRes (Certs
se_cs forall a. Semigroup a => a -> a -> a
<> Certs
cs') SubExp
se'
instance Simplifiable () where
simplify :: forall {k} (rep :: k). SimplifiableRep rep => () -> SimpleM rep ()
simplify = forall (f :: * -> *) a. Applicative f => a -> f a
pure
instance Simplifiable VName where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
VName -> SimpleM rep VName
simplify VName
v = do
Maybe (SubExp, Certs)
se <- forall {k} (rep :: k).
VName -> SymbolTable rep -> Maybe (SubExp, Certs)
ST.lookupSubExp VName
v forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (rep :: k). SimpleM rep (SymbolTable (Wise rep))
askVtable
case Maybe (SubExp, Certs)
se of
Just (Var VName
v', Certs
cs) -> do
forall {k} (rep :: k). SimpleM rep ()
changed
forall {k} (rep :: k). Certs -> SimpleM rep ()
usedCerts Certs
cs
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v'
Maybe (SubExp, Certs)
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
instance Simplifiable d => Simplifiable (ShapeBase d) where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
ShapeBase d -> SimpleM rep (ShapeBase d)
simplify = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall d. [d] -> ShapeBase d
Shape forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d. ShapeBase d -> [d]
shapeDims
instance Simplifiable ExtSize where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
ExtSize -> SimpleM rep ExtSize
simplify (Free SubExp
se) = forall a. a -> Ext a
Free forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify SubExp
se
simplify (Ext Int
x) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Int -> Ext a
Ext Int
x
instance Simplifiable Space where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
Space -> SimpleM rep Space
simplify (ScalarSpace [SubExp]
ds PrimType
t) = [SubExp] -> PrimType -> Space
ScalarSpace forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify [SubExp]
ds forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
t
simplify Space
s = forall (f :: * -> *) a. Applicative f => a -> f a
pure Space
s
instance Simplifiable PrimType where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
PrimType -> SimpleM rep PrimType
simplify = forall (f :: * -> *) a. Applicative f => a -> f a
pure
instance Simplifiable shape => Simplifiable (TypeBase shape u) where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
TypeBase shape u -> SimpleM rep (TypeBase shape u)
simplify (Array PrimType
et shape
shape u
u) =
forall shape u. PrimType -> shape -> u -> TypeBase shape u
Array forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify PrimType
et forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify shape
shape forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure u
u
simplify (Acc VName
acc Shape
ispace [Type]
ts u
u) =
forall shape u. VName -> Shape -> [Type] -> u -> TypeBase shape u
Acc forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify VName
acc forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify Shape
ispace forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify [Type]
ts forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure u
u
simplify (Mem Space
space) =
forall shape u. Space -> TypeBase shape u
Mem forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify Space
space
simplify (Prim PrimType
bt) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall shape u. PrimType -> TypeBase shape u
Prim PrimType
bt
instance Simplifiable d => Simplifiable (DimIndex d) where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
DimIndex d -> SimpleM rep (DimIndex d)
simplify (DimFix d
i) = forall d. d -> DimIndex d
DimFix forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify d
i
simplify (DimSlice d
i d
n d
s) = forall d. d -> d -> d -> DimIndex d
DimSlice forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify d
i forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify d
n forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify d
s
instance Simplifiable d => Simplifiable (Slice d) where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
Slice d -> SimpleM rep (Slice d)
simplify = forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify
simplifyLambda ::
SimplifiableRep rep =>
Lambda (Wise rep) ->
SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
simplifyLambda :: forall {k} (rep :: k).
SimplifiableRep rep =>
Lambda (Wise rep)
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
simplifyLambda Lambda (Wise rep)
lam = do
BlockPred (Wise rep)
par_blocker <- forall {k} (rep :: k) a. (Env rep -> a) -> SimpleM rep a
asksEngineEnv forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). HoistBlockers rep -> BlockPred (Wise rep)
blockHoistPar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). Env rep -> HoistBlockers rep
envHoistBlockers
forall {k} (rep :: k).
SimplifiableRep rep =>
BlockPred (Wise rep)
-> UsageTable
-> Lambda (Wise rep)
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
simplifyLambdaMaybeHoist BlockPred (Wise rep)
par_blocker forall a. Monoid a => a
mempty Lambda (Wise rep)
lam
simplifyLambdaNoHoisting ::
SimplifiableRep rep =>
Lambda (Wise rep) ->
SimpleM rep (Lambda (Wise rep))
simplifyLambdaNoHoisting :: forall {k} (rep :: k).
SimplifiableRep rep =>
Lambda (Wise rep) -> SimpleM rep (Lambda (Wise rep))
simplifyLambdaNoHoisting Lambda (Wise rep)
lam =
forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (rep :: k).
SimplifiableRep rep =>
BlockPred (Wise rep)
-> UsageTable
-> Lambda (Wise rep)
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
simplifyLambdaMaybeHoist (forall {k} (rep :: k). Bool -> BlockPred rep
isFalse Bool
False) forall a. Monoid a => a
mempty Lambda (Wise rep)
lam
simplifyLambdaMaybeHoist ::
SimplifiableRep rep =>
BlockPred (Wise rep) ->
UT.UsageTable ->
Lambda (Wise rep) ->
SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
simplifyLambdaMaybeHoist :: forall {k} (rep :: k).
SimplifiableRep rep =>
BlockPred (Wise rep)
-> UsageTable
-> Lambda (Wise rep)
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
simplifyLambdaMaybeHoist = forall {k} (rep :: k).
SimplifiableRep rep =>
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> BlockPred (Wise rep)
-> UsageTable
-> Lambda (Wise rep)
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
simplifyLambdaWith forall a. a -> a
id
simplifyLambdaWith ::
SimplifiableRep rep =>
(ST.SymbolTable (Wise rep) -> ST.SymbolTable (Wise rep)) ->
BlockPred (Wise rep) ->
UT.UsageTable ->
Lambda (Wise rep) ->
SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
simplifyLambdaWith :: forall {k} (rep :: k).
SimplifiableRep rep =>
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> BlockPred (Wise rep)
-> UsageTable
-> Lambda (Wise rep)
-> SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
simplifyLambdaWith SymbolTable (Wise rep) -> SymbolTable (Wise rep)
f BlockPred (Wise rep)
blocked UsageTable
usage lam :: Lambda (Wise rep)
lam@(Lambda [LParam (Wise rep)]
params Body (Wise rep)
body [Type]
rettype) = do
[Param (LParamInfo rep)]
params' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify) [LParam (Wise rep)]
params
let paramnames :: Names
paramnames = [VName] -> Names
namesFromList forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Lambda rep -> [VName]
boundByLambda Lambda (Wise rep)
lam
(Stms (Wise rep)
hoisted, Body (Wise rep)
body') <-
forall {k} (rep :: k) a.
SimplifiableRep rep =>
[LParam (Wise rep)] -> SimpleM rep a -> SimpleM rep a
bindLParams [Param (LParamInfo rep)]
params' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k) a.
(SymbolTable (Wise rep) -> SymbolTable (Wise rep))
-> SimpleM rep a -> SimpleM rep a
localVtable SymbolTable (Wise rep) -> SymbolTable (Wise rep)
f forall a b. (a -> b) -> a -> b
$
forall {k} (rep :: k).
SimplifiableRep rep =>
BlockPred (Wise rep)
-> UsageTable
-> [Usages]
-> Body (Wise rep)
-> SimpleM rep (Stms (Wise rep), Body (Wise rep))
simplifyBody
(BlockPred (Wise rep)
blocked forall {k} (rep :: k).
BlockPred rep -> BlockPred rep -> BlockPred rep
`orIf` forall {k} (rep :: k). ASTRep rep => Names -> BlockPred rep
hasFree Names
paramnames forall {k} (rep :: k).
BlockPred rep -> BlockPred rep -> BlockPred rep
`orIf` forall {k} (rep :: k). BlockPred rep
isConsumed)
UsageTable
usage
(forall a b. (a -> b) -> [a] -> [b]
map (forall a b. a -> b -> a
const forall a. Monoid a => a
mempty) [Type]
rettype)
Body (Wise rep)
body
[Type]
rettype' <- forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify [Type]
rettype
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall {k} (rep :: k).
[LParam rep] -> Body rep -> [Type] -> Lambda rep
Lambda [Param (LParamInfo rep)]
params' Body (Wise rep)
body' [Type]
rettype', Stms (Wise rep)
hoisted)
instance Simplifiable Certs where
simplify :: forall {k} (rep :: k).
SimplifiableRep rep =>
Certs -> SimpleM rep Certs
simplify (Certs [VName]
ocs) = [VName] -> Certs
Certs forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => [a] -> [a]
nubOrd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {k} {rep :: k}. VName -> SimpleM rep [VName]
check [VName]
ocs
where
check :: VName -> SimpleM rep [VName]
check VName
idd = do
Maybe (SubExp, Certs)
vv <- forall {k} (rep :: k).
VName -> SymbolTable rep -> Maybe (SubExp, Certs)
ST.lookupSubExp VName
idd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (rep :: k). SimpleM rep (SymbolTable (Wise rep))
askVtable
case Maybe (SubExp, Certs)
vv of
Just (Constant PrimValue
_, Certs [VName]
cs) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName]
cs
Just (Var VName
idd', Certs
_) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName
idd']
Maybe (SubExp, Certs)
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName
idd]
simplifyFun ::
SimplifiableRep rep =>
FunDef (Wise rep) ->
SimpleM rep (FunDef (Wise rep))
simplifyFun :: forall {k} (rep :: k).
SimplifiableRep rep =>
FunDef (Wise rep) -> SimpleM rep (FunDef (Wise rep))
simplifyFun (FunDef Maybe EntryPoint
entry Attrs
attrs Name
fname [RetType (Wise rep)]
rettype [FParam (Wise rep)]
params Body (Wise rep)
body) = do
[RetType rep]
rettype' <- forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify [RetType (Wise rep)]
rettype
[Param (FParamInfo rep)]
params' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall e {k} (rep :: k).
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify) [FParam (Wise rep)]
params
let usages :: [Usages]
usages = forall a b. (a -> b) -> [a] -> [b]
map (Diet -> Usages
usageFromDiet forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall shape. TypeBase shape Uniqueness -> Diet
diet forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t. DeclExtTyped t => t -> DeclExtType
declExtTypeOf) [RetType rep]
rettype'
Body (Wise rep)
body' <- forall {k} (rep :: k) a.
SimplifiableRep rep =>
[FParam (Wise rep)] -> SimpleM rep a -> SimpleM rep a
bindFParams [FParam (Wise rep)]
params forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
SimplifiableRep rep =>
UsageTable
-> [Usages] -> Body (Wise rep) -> SimpleM rep (Body (Wise rep))
simplifyBodyNoHoisting forall a. Monoid a => a
mempty [Usages]
usages Body (Wise rep)
body
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
Maybe EntryPoint
-> Attrs
-> Name
-> [RetType rep]
-> [FParam rep]
-> Body rep
-> FunDef rep
FunDef Maybe EntryPoint
entry Attrs
attrs Name
fname [RetType rep]
rettype' [Param (FParamInfo rep)]
params' Body (Wise rep)
body'