-- | This module implements a compiler pass for inlining functions,
-- then removing those that have become dead.
module Futhark.Optimise.InliningDeadFun
  ( inlineAggressively,
    inlineConservatively,
    removeDeadFunctions,
  )
where

import Control.Monad
import Control.Monad.Identity
import Control.Monad.State
import Control.Parallel.Strategies
import Data.List (partition)
import Data.Map.Strict qualified as M
import Data.Set qualified as S
import Futhark.Analysis.CallGraph
import Futhark.Analysis.SymbolTable qualified as ST
import Futhark.Builder
import Futhark.IR.SOACS
import Futhark.IR.SOACS.Simplify
  ( simpleSOACS,
    simplifyConsts,
    simplifyFun,
  )
import Futhark.Optimise.CSE
import Futhark.Optimise.Simplify.Rep (addScopeWisdom, informStms)
import Futhark.Pass
import Futhark.Transform.CopyPropagate
  ( copyPropagateInFun,
    copyPropagateInProg,
  )
import Futhark.Transform.Rename

parMapM :: MonadFreshNames m => (a -> State VNameSource b) -> [a] -> m [b]
-- The special-casing of [] is quite important here!  If 'as' is
-- empty, then we might otherwise create an empty name source below,
-- which can wreak all kinds of havoc.
parMapM :: forall (m :: * -> *) a b.
MonadFreshNames m =>
(a -> State VNameSource b) -> [a] -> m [b]
parMapM a -> State VNameSource b
_ [] = forall (f :: * -> *) a. Applicative f => a -> f a
pure []
parMapM a -> State VNameSource b
f [a]
as =
  forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource forall a b. (a -> b) -> a -> b
$ \VNameSource
src ->
    let f' :: a -> (b, VNameSource)
f' a
a = forall s a. State s a -> s -> (a, s)
runState (a -> State VNameSource b
f a
a) VNameSource
src
        ([b]
bs, [VNameSource]
srcs) = forall a b. [(a, b)] -> ([a], [b])
unzip forall a b. (a -> b) -> a -> b
$ forall b a. Strategy b -> (a -> b) -> [a] -> [b]
parMap forall a. Strategy a
rpar a -> (b, VNameSource)
f' [a]
as
     in ([b]
bs, forall a. Monoid a => [a] -> a
mconcat [VNameSource]
srcs)

-- It is more efficient to shrink the program as soon as possible,
-- rather than wait until it has balooned after full inlining.  This
-- is the inverse rate at which we perform full simplification after
-- inlining.  For the other steps we just do copy propagation.  The
-- simplification rates used have been determined heuristically and
-- are probably not optimal for any given program.
inlineFunctions ::
  MonadFreshNames m =>
  Int ->
  CallGraph ->
  S.Set Name ->
  Prog SOACS ->
  m (Prog SOACS)
inlineFunctions :: forall (m :: * -> *).
MonadFreshNames m =>
Int -> CallGraph -> Set Name -> Prog SOACS -> m (Prog SOACS)
inlineFunctions Int
simplify_rate CallGraph
cg Set Name
what_should_be_inlined Prog SOACS
prog = do
  let consts :: Stms SOACS
consts = forall rep. Prog rep -> Stms rep
progConsts Prog SOACS
prog
      funs :: [FunDef SOACS]
funs = forall rep. Prog rep -> [FunDef rep]
progFuns Prog SOACS
prog
      vtable :: SymbolTable (Wise SOACS)
vtable = forall rep. ASTRep rep => Scope rep -> SymbolTable rep
ST.fromScope (forall rep. Scope rep -> Scope (Wise rep)
addScopeWisdom (forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms SOACS
consts))

  (Stms SOACS
consts', [FunDef SOACS]
funs') <- forall {f :: * -> *}.
MonadFreshNames f =>
(Int, SymbolTable (Wise SOACS))
-> (Stms SOACS, [FunDef SOACS])
-> Set Name
-> f (Stms SOACS, [FunDef SOACS])
recurse (Int
1, SymbolTable (Wise SOACS)
vtable) (Stms SOACS
consts, [FunDef SOACS]
funs) Set Name
what_should_be_inlined
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Prog SOACS
prog {progConsts :: Stms SOACS
progConsts = Stms SOACS
consts', progFuns :: [FunDef SOACS]
progFuns = [FunDef SOACS]
funs'}
  where
    fdmap :: [FunDef rep] -> Map Name (FunDef rep)
fdmap [FunDef rep]
fds = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map forall rep. FunDef rep -> Name
funDefName [FunDef rep]
fds) [FunDef rep]
fds

    noCallsTo :: Set Name -> Name -> Bool
noCallsTo Set Name
which Name
from = forall a. Set a -> Bool
S.null forall a b. (a -> b) -> a -> b
$ Name -> CallGraph -> Set Name
allCalledBy Name
from CallGraph
cg forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` Set Name
which

    recurse :: (Int, SymbolTable (Wise SOACS))
-> (Stms SOACS, [FunDef SOACS])
-> Set Name
-> f (Stms SOACS, [FunDef SOACS])
recurse (Int
i, SymbolTable (Wise SOACS)
vtable) (Stms SOACS
consts, [FunDef SOACS]
funs) Set Name
to_inline = do
      let (Set Name
to_inline_now, Set Name
to_inline_later) =
            forall a. (a -> Bool) -> Set a -> (Set a, Set a)
S.partition (Set Name -> Name -> Bool
noCallsTo Set Name
to_inline) Set Name
to_inline
          ([FunDef SOACS]
dont_inline_in, [FunDef SOACS]
to_inline_in) =
            forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (Set Name -> Name -> Bool
noCallsTo Set Name
to_inline_now forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. FunDef rep -> Name
funDefName) [FunDef SOACS]
funs

      if forall (t :: * -> *) a. Foldable t => t a -> Bool
null Set Name
to_inline_now
        then forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms SOACS
consts, [FunDef SOACS]
funs)
        else do
          let inlinemap :: Map Name (FunDef SOACS)
inlinemap =
                forall {rep}. [FunDef rep] -> Map Name (FunDef rep)
fdmap forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
filter ((forall a. Ord a => a -> Set a -> Bool
`S.member` Set Name
to_inline_now) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. FunDef rep -> Name
funDefName) [FunDef SOACS]
dont_inline_in
          (SymbolTable (Wise SOACS)
vtable', Stms SOACS
consts') <-
            if forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Name -> CallGraph -> Bool
`calledByConsts` CallGraph
cg) Set Name
to_inline_now
              then do
                Stms SOACS
consts' <-
                  forall (m :: * -> *).
MonadFreshNames m =>
Stms SOACS -> m (Stms SOACS)
simplifyConsts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep.
(AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
Bool -> Stms rep -> Stms rep
performCSEOnStms Bool
True
                    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
MonadFreshNames m =>
Map Name (FunDef SOACS) -> Stms SOACS -> m (Stms SOACS)
inlineInStms Map Name (FunDef SOACS)
inlinemap Stms SOACS
consts
                forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall rep.
(IndexOp (Op rep), Aliased rep) =>
Stms rep -> SymbolTable rep -> SymbolTable rep
ST.insertStms (forall rep. Informing rep => Stms rep -> Stms (Wise rep)
informStms Stms SOACS
consts') forall a. Monoid a => a
mempty, Stms SOACS
consts')
              else forall (f :: * -> *) a. Applicative f => a -> f a
pure (SymbolTable (Wise SOACS)
vtable, Stms SOACS
consts)

          let simplifyFun' :: FunDef SOACS -> m (FunDef SOACS)
simplifyFun' FunDef SOACS
fd
                | Int
i forall a. Integral a => a -> a -> a
`rem` Int
simplify_rate forall a. Eq a => a -> a -> Bool
== Int
0 =
                    forall (m :: * -> *) rep.
(MonadFreshNames m, SimplifiableRep rep) =>
SimpleOps rep
-> SymbolTable (Wise rep) -> FunDef rep -> m (FunDef rep)
copyPropagateInFun SimpleOps SOACS
simpleSOACS SymbolTable (Wise SOACS)
vtable'
                      forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep.
(AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
Bool -> FunDef rep -> FunDef rep
performCSEOnFunDef Bool
True
                      forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
MonadFreshNames m =>
SymbolTable (Wise SOACS) -> FunDef SOACS -> m (FunDef SOACS)
simplifyFun SymbolTable (Wise SOACS)
vtable' FunDef SOACS
fd
                | Bool
otherwise =
                    forall (m :: * -> *) rep.
(MonadFreshNames m, SimplifiableRep rep) =>
SimpleOps rep
-> SymbolTable (Wise rep) -> FunDef rep -> m (FunDef rep)
copyPropagateInFun SimpleOps SOACS
simpleSOACS SymbolTable (Wise SOACS)
vtable' FunDef SOACS
fd

              onFun :: FunDef SOACS -> StateT VNameSource Identity (FunDef SOACS)
onFun = forall {m :: * -> *}.
MonadFreshNames m =>
FunDef SOACS -> m (FunDef SOACS)
simplifyFun' forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< forall (m :: * -> *).
MonadFreshNames m =>
Map Name (FunDef SOACS) -> FunDef SOACS -> m (FunDef SOACS)
inlineInFunDef Map Name (FunDef SOACS)
inlinemap

          [FunDef SOACS]
to_inline_in' <- forall (m :: * -> *) a b.
MonadFreshNames m =>
(a -> State VNameSource b) -> [a] -> m [b]
parMapM FunDef SOACS -> StateT VNameSource Identity (FunDef SOACS)
onFun [FunDef SOACS]
to_inline_in

          (Int, SymbolTable (Wise SOACS))
-> (Stms SOACS, [FunDef SOACS])
-> Set Name
-> f (Stms SOACS, [FunDef SOACS])
recurse
            (Int
i forall a. Num a => a -> a -> a
+ Int
1, SymbolTable (Wise SOACS)
vtable')
            (Stms SOACS
consts', [FunDef SOACS]
dont_inline_in forall a. Semigroup a => a -> a -> a
<> [FunDef SOACS]
to_inline_in')
            Set Name
to_inline_later

calledOnce :: CallGraph -> S.Set Name
calledOnce :: CallGraph -> Set Name
calledOnce = forall a. Ord a => [a] -> Set a
S.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
filter ((forall a. Eq a => a -> a -> Bool
== Int
1) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
M.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. CallGraph -> Map Name Int
numOccurences

inlineBecauseTiny :: Prog SOACS -> S.Set Name
inlineBecauseTiny :: Prog SOACS -> Set Name
inlineBecauseTiny = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall {rep}. FunDef rep -> Set Name
onFunDef forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Prog rep -> [FunDef rep]
progFuns
  where
    onFunDef :: FunDef rep -> Set Name
onFunDef FunDef rep
fd
      | (forall (t :: * -> *) a. Foldable t => t a -> Int
length (forall rep. Body rep -> Stms rep
bodyStms (forall rep. FunDef rep -> Body rep
funDefBody FunDef rep
fd)) forall a. Ord a => a -> a -> Bool
< Int
2)
          Bool -> Bool -> Bool
|| (Attr
"inline" Attr -> Attrs -> Bool
`inAttrs` forall rep. FunDef rep -> Attrs
funDefAttrs FunDef rep
fd) =
          forall a. a -> Set a
S.singleton (forall rep. FunDef rep -> Name
funDefName FunDef rep
fd)
      | Bool
otherwise = forall a. Monoid a => a
mempty

-- Conservative inlining of functions that are called just once, or
-- have #[inline] on them.
consInlineFunctions :: MonadFreshNames m => Prog SOACS -> m (Prog SOACS)
consInlineFunctions :: forall (m :: * -> *).
MonadFreshNames m =>
Prog SOACS -> m (Prog SOACS)
consInlineFunctions Prog SOACS
prog =
  forall (m :: * -> *).
MonadFreshNames m =>
Int -> CallGraph -> Set Name -> Prog SOACS -> m (Prog SOACS)
inlineFunctions Int
4 CallGraph
cg (CallGraph -> Set Name
calledOnce CallGraph
cg forall a. Semigroup a => a -> a -> a
<> Prog SOACS -> Set Name
inlineBecauseTiny Prog SOACS
prog) Prog SOACS
prog
  where
    cg :: CallGraph
cg = Prog SOACS -> CallGraph
buildCallGraph Prog SOACS
prog

-- Inline everything that is not #[noinline].
aggInlineFunctions :: MonadFreshNames m => Prog SOACS -> m (Prog SOACS)
aggInlineFunctions :: forall (m :: * -> *).
MonadFreshNames m =>
Prog SOACS -> m (Prog SOACS)
aggInlineFunctions Prog SOACS
prog =
  forall (m :: * -> *).
MonadFreshNames m =>
Int -> CallGraph -> Set Name -> Prog SOACS -> m (Prog SOACS)
inlineFunctions Int
3 CallGraph
cg (forall a. Ord a => [a] -> Set a
S.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall rep. FunDef rep -> Name
funDefName forall a b. (a -> b) -> a -> b
$ forall rep. Prog rep -> [FunDef rep]
progFuns Prog SOACS
prog) Prog SOACS
prog
  where
    cg :: CallGraph
cg = Prog SOACS -> CallGraph
buildCallGraph Prog SOACS
prog

-- | @inlineInFunDef constf fdmap caller@ inlines in @calleer@ the
-- functions in @fdmap@ that are called as @constf@. At this point the
-- preconditions are that if @fdmap@ is not empty, and, more
-- importantly, the functions in @fdmap@ do not call any other
-- functions.
inlineInFunDef ::
  MonadFreshNames m =>
  M.Map Name (FunDef SOACS) ->
  FunDef SOACS ->
  m (FunDef SOACS)
inlineInFunDef :: forall (m :: * -> *).
MonadFreshNames m =>
Map Name (FunDef SOACS) -> FunDef SOACS -> m (FunDef SOACS)
inlineInFunDef Map Name (FunDef SOACS)
fdmap (FunDef Maybe EntryPoint
entry Attrs
attrs Name
name [RetType SOACS]
rtp [FParam SOACS]
args Body SOACS
body) =
  forall rep.
Maybe EntryPoint
-> Attrs
-> Name
-> [RetType rep]
-> [FParam rep]
-> Body rep
-> FunDef rep
FunDef Maybe EntryPoint
entry Attrs
attrs Name
name [RetType SOACS]
rtp [FParam SOACS]
args forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
MonadFreshNames m =>
Map Name (FunDef SOACS) -> Body SOACS -> m (Body SOACS)
inlineInBody Map Name (FunDef SOACS)
fdmap Body SOACS
body

inlineFunction ::
  MonadFreshNames m =>
  Pat Type ->
  StmAux dec ->
  [(SubExp, Diet)] ->
  (Safety, SrcLoc, [SrcLoc]) ->
  FunDef SOACS ->
  m (Stms SOACS)
inlineFunction :: forall (m :: * -> *) dec.
MonadFreshNames m =>
Pat Type
-> StmAux dec
-> [(SubExp, Diet)]
-> (Safety, SrcLoc, [SrcLoc])
-> FunDef SOACS
-> m (Stms SOACS)
inlineFunction Pat Type
pat StmAux dec
aux [(SubExp, Diet)]
args (Safety
safety, SrcLoc
loc, [SrcLoc]
locs) FunDef SOACS
fun = do
  Body BodyDec SOACS
_ Stms SOACS
stms Result
res <-
    forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Body rep -> m (Body rep)
renameBody forall a b. (a -> b) -> a -> b
$ forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody (Stms SOACS
param_stms forall a. Semigroup a => a -> a -> a
<> Stms SOACS
body_stms) (forall rep. Body rep -> Result
bodyResult (forall rep. FunDef rep -> Body rep
funDefBody FunDef SOACS
fun))
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Stms SOACS
stms forall a. Semigroup a => a -> a -> a
<> forall rep. [Stm rep] -> Stms rep
stmsFromList (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall {rep}. Buildable rep => Ident -> SubExpRes -> Stm rep
bindSubExpRes (forall dec. Typed dec => Pat dec -> [Ident]
patIdents Pat Type
pat) Result
res)
  where
    param_stms :: Stms SOACS
param_stms =
      forall rep. [Stm rep] -> Stms rep
stmsFromList forall a b. (a -> b) -> a -> b
$
        forall rep. Certs -> Stm rep -> Stm rep
certify (forall dec. StmAux dec -> Certs
stmAuxCerts StmAux dec
aux)
          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall {rep}. Buildable rep => Ident -> SubExp -> Stm rep
bindSubExp (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Typed dec => Param dec -> Ident
paramIdent forall a b. (a -> b) -> a -> b
$ forall rep. FunDef rep -> [FParam rep]
funDefParams FunDef SOACS
fun) (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(SubExp, Diet)]
args)

    body_stms :: Stms SOACS
body_stms =
      Attrs -> Safety -> [SrcLoc] -> Stms SOACS -> Stms SOACS
addLocations (forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux dec
aux) Safety
safety (forall a. (a -> Bool) -> [a] -> [a]
filter SrcLoc -> Bool
notmempty (SrcLoc
loc forall a. a -> [a] -> [a]
: [SrcLoc]
locs)) forall a b. (a -> b) -> a -> b
$
        forall rep. Body rep -> Stms rep
bodyStms forall a b. (a -> b) -> a -> b
$
          forall rep. FunDef rep -> Body rep
funDefBody FunDef SOACS
fun

    -- Note that the sizes of arrays may not be correct at this
    -- point - it is crucial that we run copy propagation before
    -- the type checker sees this!
    bindSubExp :: Ident -> SubExp -> Stm rep
bindSubExp Ident
ident SubExp
se =
      forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet [Ident
ident] forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se

    bindSubExpRes :: Ident -> SubExpRes -> Stm rep
bindSubExpRes Ident
ident (SubExpRes Certs
cs SubExp
se) =
      forall rep. Certs -> Stm rep -> Stm rep
certify Certs
cs forall a b. (a -> b) -> a -> b
$ forall {rep}. Buildable rep => Ident -> SubExp -> Stm rep
bindSubExp Ident
ident SubExp
se

    notmempty :: SrcLoc -> Bool
notmempty = (forall a. Eq a => a -> a -> Bool
/= forall a. Monoid a => a
mempty) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Located a => a -> Loc
locOf

inlineInStms ::
  MonadFreshNames m =>
  M.Map Name (FunDef SOACS) ->
  Stms SOACS ->
  m (Stms SOACS)
inlineInStms :: forall (m :: * -> *).
MonadFreshNames m =>
Map Name (FunDef SOACS) -> Stms SOACS -> m (Stms SOACS)
inlineInStms Map Name (FunDef SOACS)
fdmap Stms SOACS
stms =
  forall rep. Body rep -> Stms rep
bodyStms forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
MonadFreshNames m =>
Map Name (FunDef SOACS) -> Body SOACS -> m (Body SOACS)
inlineInBody Map Name (FunDef SOACS)
fdmap (forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms SOACS
stms [])

inlineInBody ::
  MonadFreshNames m =>
  M.Map Name (FunDef SOACS) ->
  Body SOACS ->
  m (Body SOACS)
inlineInBody :: forall (m :: * -> *).
MonadFreshNames m =>
Map Name (FunDef SOACS) -> Body SOACS -> m (Body SOACS)
inlineInBody Map Name (FunDef SOACS)
fdmap = Body SOACS -> m (Body SOACS)
onBody
  where
    inline :: [Stm SOACS] -> m (Stms SOACS)
inline (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Apply Name
fname [(SubExp, Diet)]
args [RetType SOACS]
_ (Safety, SrcLoc, [SrcLoc])
what) : [Stm SOACS]
rest)
      | Just FunDef SOACS
fd <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fname Map Name (FunDef SOACS)
fdmap,
        Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ Attr
"noinline" Attr -> Attrs -> Bool
`inAttrs` forall rep. FunDef rep -> Attrs
funDefAttrs FunDef SOACS
fd,
        Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ Attr
"noinline" Attr -> Attrs -> Bool
`inAttrs` forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux (ExpDec SOACS)
aux =
          forall a. Semigroup a => a -> a -> a
(<>) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) dec.
MonadFreshNames m =>
Pat Type
-> StmAux dec
-> [(SubExp, Diet)]
-> (Safety, SrcLoc, [SrcLoc])
-> FunDef SOACS
-> m (Stms SOACS)
inlineFunction Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux [(SubExp, Diet)]
args (Safety, SrcLoc, [SrcLoc])
what FunDef SOACS
fd forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Stm SOACS] -> m (Stms SOACS)
inline [Stm SOACS]
rest
    inline (stm :: Stm SOACS
stm@(Let Pat (LetDec SOACS)
_ StmAux (ExpDec SOACS)
_ BasicOp {}) : [Stm SOACS]
rest) =
      (forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
stm <>) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Stm SOACS] -> m (Stms SOACS)
inline [Stm SOACS]
rest
    inline (Stm SOACS
stm : [Stm SOACS]
rest) =
      forall a. Semigroup a => a -> a -> a
(<>) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall rep. Stm rep -> Stms rep
oneStm forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Stm SOACS -> m (Stm SOACS)
onStm Stm SOACS
stm) forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Stm SOACS] -> m (Stms SOACS)
inline [Stm SOACS]
rest
    inline [] =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty

    onBody :: Body SOACS -> m (Body SOACS)
onBody (Body BodyDec SOACS
dec Stms SOACS
stms Result
res) =
      forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body BodyDec SOACS
dec forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Stm SOACS] -> m (Stms SOACS)
inline (forall rep. Stms rep -> [Stm rep]
stmsToList Stms SOACS
stms) forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure Result
res
    onStm :: Stm SOACS -> m (Stm SOACS)
onStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux Exp SOACS
e) = forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper SOACS SOACS m
inliner Exp SOACS
e

    inliner :: Mapper SOACS SOACS m
inliner =
      (forall rep (m :: * -> *). Monad m => Mapper rep rep m
identityMapper @SOACS)
        { mapOnBody :: Scope SOACS -> Body SOACS -> m (Body SOACS)
mapOnBody = forall a b. a -> b -> a
const Body SOACS -> m (Body SOACS)
onBody,
          mapOnOp :: Op SOACS -> m (Op SOACS)
mapOnOp = SOAC SOACS -> m (SOAC SOACS)
onSOAC
        }

    onSOAC :: SOAC SOACS -> m (SOAC SOACS)
onSOAC =
      forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM forall rep (m :: * -> *). Monad m => SOACMapper rep rep m
identitySOACMapper {mapOnSOACLambda :: Lambda SOACS -> m (Lambda SOACS)
mapOnSOACLambda = Lambda SOACS -> m (Lambda SOACS)
onLambda}

    onLambda :: Lambda SOACS -> m (Lambda SOACS)
onLambda (Lambda [LParam SOACS]
params Body SOACS
body [Type]
ret) =
      forall rep. [LParam rep] -> Body rep -> [Type] -> Lambda rep
Lambda [LParam SOACS]
params forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Body SOACS -> m (Body SOACS)
onBody Body SOACS
body forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure [Type]
ret

-- Propagate source locations and attributes to the inlined
-- statements.  Attributes are propagated only when applicable (this
-- probably means that every supported attribute needs to be handled
-- specially here).
addLocations :: Attrs -> Safety -> [SrcLoc] -> Stms SOACS -> Stms SOACS
addLocations :: Attrs -> Safety -> [SrcLoc] -> Stms SOACS -> Stms SOACS
addLocations Attrs
attrs Safety
caller_safety [SrcLoc]
more_locs = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Stm SOACS -> Stm SOACS
onStm
  where
    onStm :: Stm SOACS -> Stm SOACS
onStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Apply Name
fname [(SubExp, Diet)]
args [RetType SOACS]
t (Safety
safety, SrcLoc
loc, [SrcLoc]
locs))) =
      forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux' forall a b. (a -> b) -> a -> b
$
        forall rep.
Name
-> [(SubExp, Diet)]
-> [RetType rep]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp rep
Apply Name
fname [(SubExp, Diet)]
args [RetType SOACS]
t (forall a. Ord a => a -> a -> a
min Safety
caller_safety Safety
safety, SrcLoc
loc, [SrcLoc]
locs forall a. [a] -> [a] -> [a]
++ [SrcLoc]
more_locs)
      where
        aux' :: StmAux (ExpDec SOACS)
aux' = StmAux (ExpDec SOACS)
aux {stmAuxAttrs :: Attrs
stmAuxAttrs = Attrs
attrs forall a. Semigroup a => a -> a -> a
<> forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux (ExpDec SOACS)
aux}
    onStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (BasicOp (Assert SubExp
cond ErrorMsg SubExp
desc (SrcLoc
loc, [SrcLoc]
locs)))) =
      forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
pat (forall {dec}. Attrs -> StmAux dec -> StmAux dec
withAttrs (Attrs -> Attrs
attrsForAssert Attrs
attrs) StmAux (ExpDec SOACS)
aux) forall a b. (a -> b) -> a -> b
$
        case Safety
caller_safety of
          Safety
Safe -> forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> ErrorMsg SubExp -> (SrcLoc, [SrcLoc]) -> BasicOp
Assert SubExp
cond ErrorMsg SubExp
desc (SrcLoc
loc, [SrcLoc]
locs forall a. [a] -> [a] -> [a]
++ [SrcLoc]
more_locs)
          Safety
Unsafe -> forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant PrimValue
UnitValue
    onStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Op Op SOACS
soac)) =
      forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
pat (forall {dec}. Attrs -> StmAux dec -> StmAux dec
withAttrs Attrs
attrs' StmAux (ExpDec SOACS)
aux) forall a b. (a -> b) -> a -> b
$
        forall rep. Op rep -> Exp rep
Op forall a b. (a -> b) -> a -> b
$
          forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$
            forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM
              forall rep (m :: * -> *). Monad m => SOACMapper rep rep m
identitySOACMapper
                { mapOnSOACLambda :: Lambda SOACS -> Identity (Lambda SOACS)
mapOnSOACLambda = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda SOACS -> Lambda SOACS
onLambda
                }
              Op SOACS
soac
      where
        attrs' :: Attrs
attrs' = Attrs
attrs Attrs -> Attrs -> Attrs
`withoutAttrs` Attrs
for_assert
        for_assert :: Attrs
for_assert = Attrs -> Attrs
attrsForAssert Attrs
attrs
        onLambda :: Lambda SOACS -> Lambda SOACS
onLambda Lambda SOACS
lam =
          Lambda SOACS
lam {lambdaBody :: Body SOACS
lambdaBody = Attrs -> Body SOACS -> Body SOACS
onBody Attrs
for_assert forall a b. (a -> b) -> a -> b
$ forall rep. Lambda rep -> Body rep
lambdaBody Lambda SOACS
lam}
    onStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux Exp SOACS
e) =
      forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux forall a b. (a -> b) -> a -> b
$ Exp SOACS -> Exp SOACS
onExp Exp SOACS
e

    onExp :: Exp SOACS -> Exp SOACS
onExp =
      forall frep trep. Mapper frep trep Identity -> Exp frep -> Exp trep
mapExp
        forall rep (m :: * -> *). Monad m => Mapper rep rep m
identityMapper
          { mapOnBody :: Scope SOACS -> Body SOACS -> Identity (Body SOACS)
mapOnBody = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Attrs -> Body SOACS -> Body SOACS
onBody Attrs
attrs
          }

    withAttrs :: Attrs -> StmAux dec -> StmAux dec
withAttrs Attrs
attrs' StmAux dec
aux = StmAux dec
aux {stmAuxAttrs :: Attrs
stmAuxAttrs = Attrs
attrs' forall a. Semigroup a => a -> a -> a
<> forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux dec
aux}

    onBody :: Attrs -> Body SOACS -> Body SOACS
onBody Attrs
attrs' Body SOACS
body =
      Body SOACS
body
        { bodyStms :: Stms SOACS
bodyStms =
            Attrs -> Safety -> [SrcLoc] -> Stms SOACS -> Stms SOACS
addLocations Attrs
attrs' Safety
caller_safety [SrcLoc]
more_locs forall a b. (a -> b) -> a -> b
$
              forall rep. Body rep -> Stms rep
bodyStms Body SOACS
body
        }

-- | Remove functions not ultimately called from an entry point or a
-- constant.
removeDeadFunctionsF :: Prog SOACS -> Prog SOACS
removeDeadFunctionsF :: Prog SOACS -> Prog SOACS
removeDeadFunctionsF Prog SOACS
prog =
  let cg :: CallGraph
cg = Prog SOACS -> CallGraph
buildCallGraph Prog SOACS
prog
      live_funs :: [FunDef SOACS]
live_funs = forall a. (a -> Bool) -> [a] -> [a]
filter ((Name -> CallGraph -> Bool
`isFunInCallGraph` CallGraph
cg) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. FunDef rep -> Name
funDefName) forall a b. (a -> b) -> a -> b
$ forall rep. Prog rep -> [FunDef rep]
progFuns Prog SOACS
prog
   in Prog SOACS
prog {progFuns :: [FunDef SOACS]
progFuns = [FunDef SOACS]
live_funs}

-- | Inline all functions and remove the resulting dead functions.
inlineAggressively :: Pass SOACS SOACS
inlineAggressively :: Pass SOACS SOACS
inlineAggressively =
  Pass
    { passName :: String
passName = String
"Inline aggressively",
      passDescription :: String
passDescription = String
"Aggressively inline and remove resulting dead functions.",
      passFunction :: Prog SOACS -> PassM (Prog SOACS)
passFunction =
        forall rep.
SimplifiableRep rep =>
SimpleOps rep -> Prog rep -> PassM (Prog rep)
copyPropagateInProg SimpleOps SOACS
simpleSOACS forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog SOACS -> Prog SOACS
removeDeadFunctionsF forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< forall (m :: * -> *).
MonadFreshNames m =>
Prog SOACS -> m (Prog SOACS)
aggInlineFunctions
    }

-- | Inline some functions and remove the resulting dead functions.
inlineConservatively :: Pass SOACS SOACS
inlineConservatively :: Pass SOACS SOACS
inlineConservatively =
  Pass
    { passName :: String
passName = String
"Inline conservatively",
      passDescription :: String
passDescription = String
"Conservatively inline and remove resulting dead functions.",
      passFunction :: Prog SOACS -> PassM (Prog SOACS)
passFunction =
        forall rep.
SimplifiableRep rep =>
SimpleOps rep -> Prog rep -> PassM (Prog rep)
copyPropagateInProg SimpleOps SOACS
simpleSOACS forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog SOACS -> Prog SOACS
removeDeadFunctionsF forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< forall (m :: * -> *).
MonadFreshNames m =>
Prog SOACS -> m (Prog SOACS)
consInlineFunctions
    }

-- | @removeDeadFunctions prog@ removes the functions that are unreachable from
-- the main function from the program.
removeDeadFunctions :: Pass SOACS SOACS
removeDeadFunctions :: Pass SOACS SOACS
removeDeadFunctions =
  Pass
    { passName :: String
passName = String
"Remove dead functions",
      passDescription :: String
passDescription = String
"Remove the functions that are unreachable from entry points",
      passFunction :: Prog SOACS -> PassM (Prog SOACS)
passFunction = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog SOACS -> Prog SOACS
removeDeadFunctionsF
    }