{-# LANGUAGE TypeFamilies #-}

-- | Last use analysis for array short circuiting
--
-- Last-Use analysis of a Futhark program in aliased explicit-memory lore form.
-- Takes as input such a program or a function and produces a 'M.Map VName
-- Names', in which the key identified the let stmt, and the list argument
-- identifies the variables that were lastly used in that stmt.  Note that the
-- results of a body do not have a last use, and neither do a function
-- parameters if it happens to not be used inside function's body.  Such cases
-- are supposed to be treated separately.
module Futhark.Analysis.LastUse
  ( lastUseSeqMem,
    lastUseGPUMem,
    lastUseMCMem,
    LUTabFun,
    LUTabProg,
  )
where

import Control.Monad
import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.Bifunctor (bimap)
import Data.Function ((&))
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Sequence (Seq (..))
import Futhark.IR.Aliases
import Futhark.IR.GPUMem
import Futhark.IR.GPUMem qualified as GPU
import Futhark.IR.MCMem
import Futhark.IR.MCMem qualified as MC
import Futhark.IR.SeqMem
import Futhark.Optimise.ArrayShortCircuiting.DataStructs
import Futhark.Util

-- | Maps a name indentifying a Stm to the last uses in that Stm.
type LUTabFun = M.Map VName Names

-- | LU-table for the constants, and for each function.
type LUTabProg = (LUTabFun, M.Map Name LUTabFun)

type LastUseOp rep = Op (Aliases rep) -> Names -> LastUseM rep (LUTabFun, Names, Names)

-- | 'LastUseReader' allows us to abstract over representations by supplying the
-- 'onOp' function.
data LastUseReader rep = LastUseReader
  { forall rep. LastUseReader rep -> LastUseOp rep
onOp :: LastUseOp rep,
    forall rep. LastUseReader rep -> Scope (Aliases rep)
scope :: Scope (Aliases rep)
  }

-- | Maps a variable or memory block to its aliases.
type AliasTab = M.Map VName Names

newtype LastUseM rep a = LastUseM (StateT AliasTab (Reader (LastUseReader rep)) a)
  deriving
    ( Applicative (LastUseM rep)
Applicative (LastUseM rep) =>
(forall a b.
 LastUseM rep a -> (a -> LastUseM rep b) -> LastUseM rep b)
-> (forall a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep b)
-> (forall a. a -> LastUseM rep a)
-> Monad (LastUseM rep)
forall rep. Applicative (LastUseM rep)
forall a. a -> LastUseM rep a
forall rep a. a -> LastUseM rep a
forall a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep b
forall a b.
LastUseM rep a -> (a -> LastUseM rep b) -> LastUseM rep b
forall rep a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep b
forall rep a b.
LastUseM rep a -> (a -> LastUseM rep b) -> LastUseM 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
$c>>= :: forall rep a b.
LastUseM rep a -> (a -> LastUseM rep b) -> LastUseM rep b
>>= :: forall a b.
LastUseM rep a -> (a -> LastUseM rep b) -> LastUseM rep b
$c>> :: forall rep a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep b
>> :: forall a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep b
$creturn :: forall rep a. a -> LastUseM rep a
return :: forall a. a -> LastUseM rep a
Monad,
      (forall a b. (a -> b) -> LastUseM rep a -> LastUseM rep b)
-> (forall a b. a -> LastUseM rep b -> LastUseM rep a)
-> Functor (LastUseM rep)
forall a b. a -> LastUseM rep b -> LastUseM rep a
forall a b. (a -> b) -> LastUseM rep a -> LastUseM rep b
forall rep a b. a -> LastUseM rep b -> LastUseM rep a
forall rep a b. (a -> b) -> LastUseM rep a -> LastUseM rep b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall rep a b. (a -> b) -> LastUseM rep a -> LastUseM rep b
fmap :: forall a b. (a -> b) -> LastUseM rep a -> LastUseM rep b
$c<$ :: forall rep a b. a -> LastUseM rep b -> LastUseM rep a
<$ :: forall a b. a -> LastUseM rep b -> LastUseM rep a
Functor,
      Functor (LastUseM rep)
Functor (LastUseM rep) =>
(forall a. a -> LastUseM rep a)
-> (forall a b.
    LastUseM rep (a -> b) -> LastUseM rep a -> LastUseM rep b)
-> (forall a b c.
    (a -> b -> c)
    -> LastUseM rep a -> LastUseM rep b -> LastUseM rep c)
-> (forall a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep b)
-> (forall a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep a)
-> Applicative (LastUseM rep)
forall rep. Functor (LastUseM rep)
forall a. a -> LastUseM rep a
forall rep a. a -> LastUseM rep a
forall a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep a
forall a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep b
forall a b.
LastUseM rep (a -> b) -> LastUseM rep a -> LastUseM rep b
forall rep a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep a
forall rep a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep b
forall rep a b.
LastUseM rep (a -> b) -> LastUseM rep a -> LastUseM rep b
forall a b c.
(a -> b -> c) -> LastUseM rep a -> LastUseM rep b -> LastUseM rep c
forall rep a b c.
(a -> b -> c) -> LastUseM rep a -> LastUseM rep b -> LastUseM 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
$cpure :: forall rep a. a -> LastUseM rep a
pure :: forall a. a -> LastUseM rep a
$c<*> :: forall rep a b.
LastUseM rep (a -> b) -> LastUseM rep a -> LastUseM rep b
<*> :: forall a b.
LastUseM rep (a -> b) -> LastUseM rep a -> LastUseM rep b
$cliftA2 :: forall rep a b c.
(a -> b -> c) -> LastUseM rep a -> LastUseM rep b -> LastUseM rep c
liftA2 :: forall a b c.
(a -> b -> c) -> LastUseM rep a -> LastUseM rep b -> LastUseM rep c
$c*> :: forall rep a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep b
*> :: forall a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep b
$c<* :: forall rep a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep a
<* :: forall a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep a
Applicative,
      MonadReader (LastUseReader rep),
      MonadState AliasTab
    )

instance (RepTypes (Aliases rep)) => HasScope (Aliases rep) (LastUseM rep) where
  askScope :: LastUseM rep (Scope (Aliases rep))
askScope = (LastUseReader rep -> Scope (Aliases rep))
-> LastUseM rep (Scope (Aliases rep))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks LastUseReader rep -> Scope (Aliases rep)
forall rep. LastUseReader rep -> Scope (Aliases rep)
scope

instance (RepTypes (Aliases rep)) => LocalScope (Aliases rep) (LastUseM rep) where
  localScope :: forall a. Scope (Aliases rep) -> LastUseM rep a -> LastUseM rep a
localScope Scope (Aliases rep)
sc (LastUseM StateT LUTabFun (Reader (LastUseReader rep)) a
m) = StateT LUTabFun (Reader (LastUseReader rep)) a -> LastUseM rep a
forall rep a.
StateT LUTabFun (Reader (LastUseReader rep)) a -> LastUseM rep a
LastUseM (StateT LUTabFun (Reader (LastUseReader rep)) a -> LastUseM rep a)
-> StateT LUTabFun (Reader (LastUseReader rep)) a -> LastUseM rep a
forall a b. (a -> b) -> a -> b
$ do
    (LastUseReader rep -> LastUseReader rep)
-> StateT LUTabFun (Reader (LastUseReader rep)) a
-> StateT LUTabFun (Reader (LastUseReader rep)) a
forall a.
(LastUseReader rep -> LastUseReader rep)
-> StateT LUTabFun (Reader (LastUseReader rep)) a
-> StateT LUTabFun (Reader (LastUseReader rep)) a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (\LastUseReader rep
rd -> LastUseReader rep
rd {scope = scope rd <> sc}) StateT LUTabFun (Reader (LastUseReader rep)) a
m

type Constraints rep =
  ( LocalScope (Aliases rep) (LastUseM rep),
    HasMemBlock (Aliases rep),
    AliasableRep rep
  )

runLastUseM :: LastUseOp rep -> LastUseM rep a -> a
runLastUseM :: forall rep a. LastUseOp rep -> LastUseM rep a -> a
runLastUseM LastUseOp rep
onOp (LastUseM StateT LUTabFun (ReaderT (LastUseReader rep) Identity) a
m) =
  Reader (LastUseReader rep) a -> LastUseReader rep -> a
forall r a. Reader r a -> r -> a
runReader (StateT LUTabFun (ReaderT (LastUseReader rep) Identity) a
-> LUTabFun -> Reader (LastUseReader rep) a
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT StateT LUTabFun (ReaderT (LastUseReader rep) Identity) a
m LUTabFun
forall a. Monoid a => a
mempty) (LastUseOp rep -> Scope (Aliases rep) -> LastUseReader rep
forall rep.
LastUseOp rep -> Scope (Aliases rep) -> LastUseReader rep
LastUseReader LastUseOp rep
onOp Scope (Aliases rep)
forall a. Monoid a => a
mempty)

aliasLookup :: VName -> LastUseM rep Names
aliasLookup :: forall rep. VName -> LastUseM rep Names
aliasLookup VName
vname =
  (LUTabFun -> Names) -> LastUseM rep Names
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((LUTabFun -> Names) -> LastUseM rep Names)
-> (LUTabFun -> Names) -> LastUseM rep Names
forall a b. (a -> b) -> a -> b
$ Names -> Maybe Names -> Names
forall a. a -> Maybe a -> a
fromMaybe Names
forall a. Monoid a => a
mempty (Maybe Names -> Names)
-> (LUTabFun -> Maybe Names) -> LUTabFun -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> LUTabFun -> Maybe Names
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
vname

lastUseProg ::
  (Constraints rep) =>
  Prog (Aliases rep) ->
  LastUseM rep LUTabProg
lastUseProg :: forall rep.
Constraints rep =>
Prog (Aliases rep) -> LastUseM rep LUTabProg
lastUseProg Prog (Aliases rep)
prog =
  let bound_in_consts :: Names
bound_in_consts =
        Prog (Aliases rep) -> Stms (Aliases rep)
forall rep. Prog rep -> Stms rep
progConsts Prog (Aliases rep)
prog
          Stms (Aliases rep) -> (Stms (Aliases rep) -> [VName]) -> [VName]
forall a b. a -> (a -> b) -> b
& (Stm (Aliases rep) -> [VName]) -> Stms (Aliases rep) -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Pat (VarAliases, LetDec rep) -> [VName]
forall dec. Pat dec -> [VName]
patNames (Pat (VarAliases, LetDec rep) -> [VName])
-> (Stm (Aliases rep) -> Pat (VarAliases, LetDec rep))
-> Stm (Aliases rep)
-> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm (Aliases rep) -> Pat (VarAliases, LetDec rep)
Stm (Aliases rep) -> Pat (LetDec (Aliases rep))
forall rep. Stm rep -> Pat (LetDec rep)
stmPat)
          [VName] -> ([VName] -> Names) -> Names
forall a b. a -> (a -> b) -> b
& [VName] -> Names
namesFromList
      consts :: Stms (Aliases rep)
consts = Prog (Aliases rep) -> Stms (Aliases rep)
forall rep. Prog rep -> Stms rep
progConsts Prog (Aliases rep)
prog
      funs :: [FunDef (Aliases rep)]
funs = Prog (Aliases rep) -> [FunDef (Aliases rep)]
forall rep. Prog rep -> [FunDef rep]
progFuns Prog (Aliases rep)
prog
   in Stms (Aliases rep)
-> LastUseM rep LUTabProg -> LastUseM rep LUTabProg
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Stms (Aliases rep)
consts (LastUseM rep LUTabProg -> LastUseM rep LUTabProg)
-> LastUseM rep LUTabProg -> LastUseM rep LUTabProg
forall a b. (a -> b) -> a -> b
$ do
        (LUTabFun
consts_lu, Names
_) <- Stms (Aliases rep)
-> (LUTabFun, Names) -> [VName] -> LastUseM rep (LUTabFun, Names)
forall rep.
Constraints rep =>
Stms (Aliases rep)
-> (LUTabFun, Names) -> [VName] -> LastUseM rep (LUTabFun, Names)
lastUseStms Stms (Aliases rep)
consts (LUTabFun, Names)
forall a. Monoid a => a
mempty [VName]
forall a. Monoid a => a
mempty
        [LUTabFun]
lus <- (FunDef (Aliases rep) -> LastUseM rep LUTabFun)
-> [FunDef (Aliases rep)] -> LastUseM rep [LUTabFun]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Names -> FunDef (Aliases rep) -> LastUseM rep LUTabFun
forall rep.
Constraints rep =>
Names -> FunDef (Aliases rep) -> LastUseM rep LUTabFun
lastUseFun Names
bound_in_consts) [FunDef (Aliases rep)]
funs
        LUTabProg -> LastUseM rep LUTabProg
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
consts_lu, [(Name, LUTabFun)] -> Map Name LUTabFun
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, LUTabFun)] -> Map Name LUTabFun)
-> [(Name, LUTabFun)] -> Map Name LUTabFun
forall a b. (a -> b) -> a -> b
$ [Name] -> [LUTabFun] -> [(Name, LUTabFun)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((FunDef (Aliases rep) -> Name) -> [FunDef (Aliases rep)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map FunDef (Aliases rep) -> Name
forall rep. FunDef rep -> Name
funDefName [FunDef (Aliases rep)]
funs) [LUTabFun]
lus)

lastUseFun ::
  (Constraints rep) =>
  Names ->
  FunDef (Aliases rep) ->
  LastUseM rep LUTabFun
lastUseFun :: forall rep.
Constraints rep =>
Names -> FunDef (Aliases rep) -> LastUseM rep LUTabFun
lastUseFun Names
bound_in_consts FunDef (Aliases rep)
f =
  FunDef (Aliases rep)
-> LastUseM rep LUTabFun -> LastUseM rep LUTabFun
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf FunDef (Aliases rep)
f (LastUseM rep LUTabFun -> LastUseM rep LUTabFun)
-> LastUseM rep LUTabFun -> LastUseM rep LUTabFun
forall a b. (a -> b) -> a -> b
$ (LUTabFun, Names) -> LUTabFun
forall a b. (a, b) -> a
fst ((LUTabFun, Names) -> LUTabFun)
-> LastUseM rep (LUTabFun, Names) -> LastUseM rep LUTabFun
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall rep.
Constraints rep =>
Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseBody (FunDef (Aliases rep) -> Body (Aliases rep)
forall rep. FunDef rep -> Body rep
funDefBody FunDef (Aliases rep)
f) (LUTabFun
forall a. Monoid a => a
mempty, Names
bound_in_consts)

-- | Perform last-use analysis on a 'Prog' in 'SeqMem'
lastUseSeqMem :: Prog (Aliases SeqMem) -> LUTabProg
lastUseSeqMem :: Prog (Aliases SeqMem) -> LUTabProg
lastUseSeqMem = LastUseOp SeqMem -> LastUseM SeqMem LUTabProg -> LUTabProg
forall rep a. LastUseOp rep -> LastUseM rep a -> a
runLastUseM LastUseOp SeqMem
lastUseSeqOp (LastUseM SeqMem LUTabProg -> LUTabProg)
-> (Prog (Aliases SeqMem) -> LastUseM SeqMem LUTabProg)
-> Prog (Aliases SeqMem)
-> LUTabProg
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog (Aliases SeqMem) -> LastUseM SeqMem LUTabProg
forall rep.
Constraints rep =>
Prog (Aliases rep) -> LastUseM rep LUTabProg
lastUseProg

-- | Perform last-use analysis on a 'Prog' in 'GPUMem'
lastUseGPUMem :: Prog (Aliases GPUMem) -> LUTabProg
lastUseGPUMem :: Prog (Aliases GPUMem) -> LUTabProg
lastUseGPUMem = LastUseOp GPUMem -> LastUseM GPUMem LUTabProg -> LUTabProg
forall rep a. LastUseOp rep -> LastUseM rep a -> a
runLastUseM ((HostOp NoOp (Aliases GPUMem)
 -> Names -> LastUseM GPUMem (LUTabFun, Names, Names))
-> MemOp (HostOp NoOp) (Aliases GPUMem)
-> Names
-> LastUseM GPUMem (LUTabFun, Names, Names)
forall (inner :: * -> *) rep.
(inner (Aliases rep)
 -> Names -> LastUseM rep (LUTabFun, Names, Names))
-> MemOp inner (Aliases rep)
-> Names
-> LastUseM rep (LUTabFun, Names, Names)
lastUseMemOp HostOp NoOp (Aliases GPUMem)
-> Names -> LastUseM GPUMem (LUTabFun, Names, Names)
lastUseGPUOp) (LastUseM GPUMem LUTabProg -> LUTabProg)
-> (Prog (Aliases GPUMem) -> LastUseM GPUMem LUTabProg)
-> Prog (Aliases GPUMem)
-> LUTabProg
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog (Aliases GPUMem) -> LastUseM GPUMem LUTabProg
forall rep.
Constraints rep =>
Prog (Aliases rep) -> LastUseM rep LUTabProg
lastUseProg

-- | Perform last-use analysis on a 'Prog' in 'MCMem'
lastUseMCMem :: Prog (Aliases MCMem) -> LUTabProg
lastUseMCMem :: Prog (Aliases MCMem) -> LUTabProg
lastUseMCMem = LastUseOp MCMem -> LastUseM MCMem LUTabProg -> LUTabProg
forall rep a. LastUseOp rep -> LastUseM rep a -> a
runLastUseM ((MCOp NoOp (Aliases MCMem)
 -> Names -> LastUseM MCMem (LUTabFun, Names, Names))
-> MemOp (MCOp NoOp) (Aliases MCMem)
-> Names
-> LastUseM MCMem (LUTabFun, Names, Names)
forall (inner :: * -> *) rep.
(inner (Aliases rep)
 -> Names -> LastUseM rep (LUTabFun, Names, Names))
-> MemOp inner (Aliases rep)
-> Names
-> LastUseM rep (LUTabFun, Names, Names)
lastUseMemOp MCOp NoOp (Aliases MCMem)
-> Names -> LastUseM MCMem (LUTabFun, Names, Names)
lastUseMCOp) (LastUseM MCMem LUTabProg -> LUTabProg)
-> (Prog (Aliases MCMem) -> LastUseM MCMem LUTabProg)
-> Prog (Aliases MCMem)
-> LUTabProg
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog (Aliases MCMem) -> LastUseM MCMem LUTabProg
forall rep.
Constraints rep =>
Prog (Aliases rep) -> LastUseM rep LUTabProg
lastUseProg

-- | Performing the last-use analysis on a body.
--
-- The implementation consists of a bottom-up traversal of the body's statements
-- in which the the variables lastly used in a statement are computed as the
-- difference between the free-variables in that stmt and the set of variables
-- known to be used after that statement.
lastUseBody ::
  (Constraints rep) =>
  -- | The body of statements
  Body (Aliases rep) ->
  -- | The current last-use table, tupled with the known set of already used names
  (LUTabFun, Names) ->
  -- | The result is:
  --      (i) an updated last-use table,
  --     (ii) an updated set of used names (including the binding).
  LastUseM rep (LUTabFun, Names)
lastUseBody :: forall rep.
Constraints rep =>
Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseBody bdy :: Body (Aliases rep)
bdy@(Body BodyDec (Aliases rep)
_ Stms (Aliases rep)
stms Result
result) (LUTabFun
lutab, Names
used_nms) =
  -- perform analysis bottom-up in bindings: results are known to be used,
  -- hence they are added to the used_nms set.
  Stms (Aliases rep)
-> LastUseM rep (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Stms (Aliases rep)
stms (LastUseM rep (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names))
-> LastUseM rep (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall a b. (a -> b) -> a -> b
$ do
    (LUTabFun
lutab', Names
_) <-
      Stms (Aliases rep)
-> (LUTabFun, Names) -> [VName] -> LastUseM rep (LUTabFun, Names)
forall rep.
Constraints rep =>
Stms (Aliases rep)
-> (LUTabFun, Names) -> [VName] -> LastUseM rep (LUTabFun, Names)
lastUseStms Stms (Aliases rep)
stms (LUTabFun
lutab, Names
used_nms) ([VName] -> LastUseM rep (LUTabFun, Names))
-> [VName] -> LastUseM rep (LUTabFun, Names)
forall a b. (a -> b) -> a -> b
$
        Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$
          [SubExp] -> Names
forall a. FreeIn a => a -> Names
freeIn ([SubExp] -> Names) -> [SubExp] -> Names
forall a b. (a -> b) -> a -> b
$
            (SubExpRes -> SubExp) -> Result -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp Result
result
    -- Clean up the used names by recomputing the aliasing transitive-closure
    -- of the free names in body based on the current alias table @alstab@.
    Names
used_in_body <- Names -> LastUseM rep Names
forall rep. Names -> LastUseM rep Names
aliasTransitiveClosure (Names -> LastUseM rep Names) -> Names -> LastUseM rep Names
forall a b. (a -> b) -> a -> b
$ Body (Aliases rep) -> Names
forall a. FreeIn a => a -> Names
freeIn Body (Aliases rep)
bdy
    (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
lutab', Names
used_nms Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
used_in_body)

-- | Performing the last-use analysis on a body.
--
-- The implementation consists of a bottom-up traversal of the body's statements
-- in which the the variables lastly used in a statement are computed as the
-- difference between the free-variables in that stmt and the set of variables
-- known to be used after that statement.
lastUseKernelBody ::
  (Constraints rep) =>
  -- | The body of statements
  KernelBody (Aliases rep) ->
  -- | The current last-use table, tupled with the known set of already used names
  (LUTabFun, Names) ->
  -- | The result is:
  --      (i) an updated last-use table,
  --     (ii) an updated set of used names (including the binding).
  LastUseM rep (LUTabFun, Names)
lastUseKernelBody :: forall rep.
Constraints rep =>
KernelBody (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseKernelBody bdy :: KernelBody (Aliases rep)
bdy@(KernelBody BodyDec (Aliases rep)
_ Stms (Aliases rep)
stms [KernelResult]
result) (LUTabFun
lutab, Names
used_nms) =
  Stms (Aliases rep)
-> LastUseM rep (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Stms (Aliases rep)
stms (LastUseM rep (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names))
-> LastUseM rep (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall a b. (a -> b) -> a -> b
$ do
    -- perform analysis bottom-up in bindings: results are known to be used,
    -- hence they are added to the used_nms set.
    (LUTabFun
lutab', Names
_) <-
      Stms (Aliases rep)
-> (LUTabFun, Names) -> [VName] -> LastUseM rep (LUTabFun, Names)
forall rep.
Constraints rep =>
Stms (Aliases rep)
-> (LUTabFun, Names) -> [VName] -> LastUseM rep (LUTabFun, Names)
lastUseStms Stms (Aliases rep)
stms (LUTabFun
lutab, Names
used_nms) ([VName] -> LastUseM rep (LUTabFun, Names))
-> [VName] -> LastUseM rep (LUTabFun, Names)
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ [KernelResult] -> Names
forall a. FreeIn a => a -> Names
freeIn [KernelResult]
result
    -- Clean up the used names by recomputing the aliasing transitive-closure
    -- of the free names in body based on the current alias table @alstab@.
    Names
used_in_body <- Names -> LastUseM rep Names
forall rep. Names -> LastUseM rep Names
aliasTransitiveClosure (Names -> LastUseM rep Names) -> Names -> LastUseM rep Names
forall a b. (a -> b) -> a -> b
$ KernelBody (Aliases rep) -> Names
forall a. FreeIn a => a -> Names
freeIn KernelBody (Aliases rep)
bdy
    (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
lutab', Names
used_nms Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
used_in_body)

lastUseStms ::
  (Constraints rep) =>
  Stms (Aliases rep) ->
  (LUTabFun, Names) ->
  [VName] ->
  LastUseM rep (LUTabFun, Names)
lastUseStms :: forall rep.
Constraints rep =>
Stms (Aliases rep)
-> (LUTabFun, Names) -> [VName] -> LastUseM rep (LUTabFun, Names)
lastUseStms Seq (Stm (Aliases rep))
Empty (LUTabFun
lutab, Names
nms) [VName]
res_nms = do
  Names
aliases <- (VName -> LastUseM rep Names) -> [VName] -> LastUseM rep Names
forall (m :: * -> *) b a.
(Monad m, Monoid b) =>
(a -> m b) -> [a] -> m b
concatMapM VName -> LastUseM rep Names
forall rep. VName -> LastUseM rep Names
aliasLookup [VName]
res_nms
  (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
lutab, Names
nms Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
aliases Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [VName] -> Names
namesFromList [VName]
res_nms)
lastUseStms (stm :: Stm (Aliases rep)
stm@(Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
_ Exp (Aliases rep)
e) :<| Seq (Stm (Aliases rep))
stms) (LUTabFun
lutab, Names
nms) [VName]
res_nms =
  Stm (Aliases rep)
-> LastUseM rep (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Stm (Aliases rep)
stm (LastUseM rep (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names))
-> LastUseM rep (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall a b. (a -> b) -> a -> b
$ do
    let extra_alias :: Names
extra_alias = case Exp (Aliases rep)
e of
          BasicOp (Update Safety
_ VName
old Slice SubExp
_ SubExp
_) -> VName -> Names
oneName VName
old
          BasicOp (FlatUpdate VName
old FlatSlice SubExp
_ VName
_) -> VName -> Names
oneName VName
old
          Exp (Aliases rep)
_ -> Names
forall a. Monoid a => a
mempty
    -- We build up aliases top-down
    Names -> Pat (VarAliases, LetDec rep) -> LastUseM rep ()
forall dec rep.
AliasesOf dec =>
Names -> Pat dec -> LastUseM rep ()
updateAliasing Names
extra_alias Pat (VarAliases, LetDec rep)
Pat (LetDec (Aliases rep))
pat
    -- But compute last use bottom-up
    (LUTabFun
lutab', Names
nms') <- Seq (Stm (Aliases rep))
-> (LUTabFun, Names) -> [VName] -> LastUseM rep (LUTabFun, Names)
forall rep.
Constraints rep =>
Stms (Aliases rep)
-> (LUTabFun, Names) -> [VName] -> LastUseM rep (LUTabFun, Names)
lastUseStms Seq (Stm (Aliases rep))
stms (LUTabFun
lutab, Names
nms) [VName]
res_nms
    (LUTabFun
lutab'', Names
nms'') <- Stm (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall rep.
Constraints rep =>
Stm (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseStm Stm (Aliases rep)
stm (LUTabFun
lutab', Names
nms')
    (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
lutab'', Names
nms'')

lastUseStm ::
  (Constraints rep) =>
  Stm (Aliases rep) ->
  (LUTabFun, Names) ->
  LastUseM rep (LUTabFun, Names)
lastUseStm :: forall rep.
Constraints rep =>
Stm (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseStm (Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
_ Exp (Aliases rep)
e) (LUTabFun
lutab, Names
used_nms) = do
  -- analyse the expression and get the
  --  (i)  a new last-use table (in case the @e@ contains bodies of stmts)
  -- (ii) the set of variables lastly used in the current binding.
  -- (iii)  aliased transitive-closure of used names, and
  (LUTabFun
lutab', Names
last_uses, Names
used_nms') <- Exp (Aliases rep) -> Names -> LastUseM rep (LUTabFun, Names, Names)
forall rep.
Constraints rep =>
Exp (Aliases rep) -> Names -> LastUseM rep (LUTabFun, Names, Names)
lastUseExp Exp (Aliases rep)
e Names
used_nms
  Scope (Aliases rep)
sc <- (LastUseReader rep -> Scope (Aliases rep))
-> LastUseM rep (Scope (Aliases rep))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks LastUseReader rep -> Scope (Aliases rep)
forall rep. LastUseReader rep -> Scope (Aliases rep)
scope
  let lu_mems :: Names
lu_mems =
        Names -> [VName]
namesToList Names
last_uses
          [VName] -> ([VName] -> [ArrayMemBound]) -> [ArrayMemBound]
forall a b. a -> (a -> b) -> b
& (VName -> Maybe ArrayMemBound) -> [VName] -> [ArrayMemBound]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (VName -> Scope (Aliases rep) -> Maybe ArrayMemBound
forall rep.
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
`getScopeMemInfo` Scope (Aliases rep)
sc)
          [ArrayMemBound] -> ([ArrayMemBound] -> [VName]) -> [VName]
forall a b. a -> (a -> b) -> b
& (ArrayMemBound -> VName) -> [ArrayMemBound] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map ArrayMemBound -> VName
memName
          [VName] -> ([VName] -> Names) -> Names
forall a b. a -> (a -> b) -> b
& [VName] -> Names
namesFromList
          Names -> (Names -> Names) -> Names
forall a b. a -> (a -> b) -> b
& (Names -> Names -> Names) -> Names -> Names -> Names
forall a b c. (a -> b -> c) -> b -> a -> c
flip Names -> Names -> Names
namesSubtract Names
used_nms

  -- filter-out the binded names from the set of used variables,
  -- since they go out of scope, and update the last-use table.
  let patnms :: [VName]
patnms = Pat (VarAliases, LetDec rep) -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat (VarAliases, LetDec rep)
Pat (LetDec (Aliases rep))
pat
      used_nms'' :: Names
used_nms'' = Names
used_nms' Names -> Names -> Names
`namesSubtract` [VName] -> Names
namesFromList [VName]
patnms
      lutab'' :: LUTabFun
lutab'' =
        LUTabFun -> LUTabFun -> LUTabFun
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union LUTabFun
lutab' (LUTabFun -> LUTabFun) -> LUTabFun -> LUTabFun
forall a b. (a -> b) -> a -> b
$ VName -> Names -> LUTabFun -> LUTabFun
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert ([VName] -> VName
forall a. HasCallStack => [a] -> a
head [VName]
patnms) (Names
last_uses Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
lu_mems) LUTabFun
lutab
  (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
lutab'', Names
used_nms'')

--------------------------------

-- | Last-Use Analysis for an expression.
lastUseExp ::
  (Constraints rep) =>
  -- | The expression to analyse
  Exp (Aliases rep) ->
  -- | The set of used names "after" this expression
  Names ->
  -- | Result:
  --    1. an extra LUTab recording the last use for expression's inner bodies,
  --    2. the set of last-used vars in the expression at this level,
  --    3. the updated used names, now including expression's free vars.
  LastUseM rep (LUTabFun, Names, Names)
lastUseExp :: forall rep.
Constraints rep =>
Exp (Aliases rep) -> Names -> LastUseM rep (LUTabFun, Names, Names)
lastUseExp (Match [SubExp]
_ [Case (Body (Aliases rep))]
cases Body (Aliases rep)
body MatchDec (BranchType (Aliases rep))
_) Names
used_nms = do
  -- For an if-then-else, we duplicate the last use at each body level, meaning
  -- we record the last use of the outer statement, and also the last use in the
  -- statement in the inner bodies. We can safely ignore the if-condition as it is
  -- a boolean scalar.
  (LUTabFun
lutab_cases, Names
used_cases) <-
    ([LUTabFun] -> LUTabFun)
-> ([Names] -> Names) -> ([LUTabFun], [Names]) -> (LUTabFun, Names)
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap [LUTabFun] -> LUTabFun
forall a. Monoid a => [a] -> a
mconcat [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat (([LUTabFun], [Names]) -> (LUTabFun, Names))
-> ([(LUTabFun, Names)] -> ([LUTabFun], [Names]))
-> [(LUTabFun, Names)]
-> (LUTabFun, Names)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(LUTabFun, Names)] -> ([LUTabFun], [Names])
forall a b. [(a, b)] -> ([a], [b])
unzip
      ([(LUTabFun, Names)] -> (LUTabFun, Names))
-> LastUseM rep [(LUTabFun, Names)]
-> LastUseM rep (LUTabFun, Names)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Case (Body (Aliases rep)) -> LastUseM rep (LUTabFun, Names))
-> [Case (Body (Aliases rep))] -> LastUseM rep [(LUTabFun, Names)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((Body (Aliases rep)
 -> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names))
-> (LUTabFun, Names)
-> Body (Aliases rep)
-> LastUseM rep (LUTabFun, Names)
forall a b c. (a -> b -> c) -> b -> a -> c
flip Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall rep.
Constraints rep =>
Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseBody (LUTabFun
forall k a. Map k a
M.empty, Names
used_nms) (Body (Aliases rep) -> LastUseM rep (LUTabFun, Names))
-> (Case (Body (Aliases rep)) -> Body (Aliases rep))
-> Case (Body (Aliases rep))
-> LastUseM rep (LUTabFun, Names)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Case (Body (Aliases rep)) -> Body (Aliases rep)
forall body. Case body -> body
caseBody) [Case (Body (Aliases rep))]
cases
  (LUTabFun
lutab', Names
body_used_nms) <- Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall rep.
Constraints rep =>
Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseBody Body (Aliases rep)
body (LUTabFun
forall k a. Map k a
M.empty, Names
used_nms)
  let free_in_body :: Names
free_in_body = Body (Aliases rep) -> Names
forall a. FreeIn a => a -> Names
freeIn Body (Aliases rep)
body
  let free_in_cases :: Names
free_in_cases = [Case (Body (Aliases rep))] -> Names
forall a. FreeIn a => a -> Names
freeIn [Case (Body (Aliases rep))]
cases
  let used_nms' :: Names
used_nms' = Names
used_cases Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
body_used_nms
  (Names
_, Names
last_used_arrs) <- Names -> Names -> LastUseM rep (Names, Names)
forall rep. Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms (Names -> LastUseM rep (Names, Names))
-> Names -> LastUseM rep (Names, Names)
forall a b. (a -> b) -> a -> b
$ Names
free_in_body Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
free_in_cases
  (LUTabFun, Names, Names) -> LastUseM rep (LUTabFun, Names, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
lutab_cases LUTabFun -> LUTabFun -> LUTabFun
forall a. Semigroup a => a -> a -> a
<> LUTabFun
lutab', Names
last_used_arrs, Names
used_nms')
lastUseExp (Loop [(FParam (Aliases rep), SubExp)]
var_ses LoopForm
form Body (Aliases rep)
body) Names
used_nms0 = Scope (Aliases rep)
-> LastUseM rep (LUTabFun, Names, Names)
-> LastUseM rep (LUTabFun, Names, Names)
forall a. Scope (Aliases rep) -> LastUseM rep a -> LastUseM rep a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (LoopForm -> Scope (Aliases rep)
forall rep. LoopForm -> Scope rep
scopeOfLoopForm LoopForm
form) (LastUseM rep (LUTabFun, Names, Names)
 -> LastUseM rep (LUTabFun, Names, Names))
-> LastUseM rep (LUTabFun, Names, Names)
-> LastUseM rep (LUTabFun, Names, Names)
forall a b. (a -> b) -> a -> b
$ do
  Names
free_in_body <- Names -> LastUseM rep Names
forall rep. Names -> LastUseM rep Names
aliasTransitiveClosure (Names -> LastUseM rep Names) -> Names -> LastUseM rep Names
forall a b. (a -> b) -> a -> b
$ Body (Aliases rep) -> Names
forall a. FreeIn a => a -> Names
freeIn Body (Aliases rep)
body
  -- compute the aliasing transitive closure of initializers that are not last-uses
  [(VName, Names)]
var_inis <- [Maybe (VName, Names)] -> [(VName, Names)]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (VName, Names)] -> [(VName, Names)])
-> LastUseM rep [Maybe (VName, Names)]
-> LastUseM rep [(VName, Names)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Param (FParamInfo rep), SubExp)
 -> LastUseM rep (Maybe (VName, Names)))
-> [(Param (FParamInfo rep), SubExp)]
-> LastUseM rep [Maybe (VName, Names)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Names
-> (Param (FParamInfo rep), SubExp)
-> LastUseM rep (Maybe (VName, Names))
forall {dec} {rep}.
Typed dec =>
Names -> (Param dec, SubExp) -> LastUseM rep (Maybe (VName, Names))
initHelper (Names
free_in_body Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
used_nms0)) [(Param (FParamInfo rep), SubExp)]
[(FParam (Aliases rep), SubExp)]
var_ses
  let -- To record last-uses inside the loop body, we call 'lastUseBody' with used-names
      -- being:  (free_in_body - loop-variants-a) + used_nms0. As such we disable cases b)
      -- and c) to produce loop-variant last uses inside the loop, and also we prevent
      -- the free-loop-variables to having last uses inside the loop.
      free_in_body' :: Names
free_in_body' = Names
free_in_body Names -> Names -> Names
`namesSubtract` [VName] -> Names
namesFromList (((VName, Names) -> VName) -> [(VName, Names)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, Names) -> VName
forall a b. (a, b) -> a
fst [(VName, Names)]
var_inis)
      used_nms :: Names
used_nms = Names
used_nms0 Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
free_in_body' Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Result -> Names
forall a. FreeIn a => a -> Names
freeIn (Body (Aliases rep) -> Result
forall rep. Body rep -> Result
bodyResult Body (Aliases rep)
body)
  (LUTabFun
body_lutab, Names
_) <- Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall rep.
Constraints rep =>
Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseBody Body (Aliases rep)
body (LUTabFun
forall a. Monoid a => a
mempty, Names
used_nms)

  -- add var_inis_a to the body_lutab, i.e., record the last-use of
  -- initializer in the corresponding loop variant.
  let lutab_res :: LUTabFun
lutab_res = LUTabFun
body_lutab LUTabFun -> LUTabFun -> LUTabFun
forall a. Semigroup a => a -> a -> a
<> [(VName, Names)] -> LUTabFun
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName, Names)]
var_inis

      -- the result used names are:
      fpar_nms :: Names
fpar_nms = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ ((Param (FParamInfo rep), SubExp) -> VName)
-> [(Param (FParamInfo rep), SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (Ident -> VName
identName (Ident -> VName)
-> ((Param (FParamInfo rep), SubExp) -> Ident)
-> (Param (FParamInfo rep), SubExp)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param (FParamInfo rep) -> Ident
forall dec. Typed dec => Param dec -> Ident
paramIdent (Param (FParamInfo rep) -> Ident)
-> ((Param (FParamInfo rep), SubExp) -> Param (FParamInfo rep))
-> (Param (FParamInfo rep), SubExp)
-> Ident
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param (FParamInfo rep), SubExp) -> Param (FParamInfo rep)
forall a b. (a, b) -> a
fst) [(Param (FParamInfo rep), SubExp)]
[(FParam (Aliases rep), SubExp)]
var_ses
      used_nms' :: Names
used_nms' = (Names
free_in_body Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> Names
forall a. FreeIn a => a -> Names
freeIn (((Param (FParamInfo rep), SubExp) -> SubExp)
-> [(Param (FParamInfo rep), SubExp)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (Param (FParamInfo rep), SubExp) -> SubExp
forall a b. (a, b) -> b
snd [(Param (FParamInfo rep), SubExp)]
[(FParam (Aliases rep), SubExp)]
var_ses)) Names -> Names -> Names
`namesSubtract` Names
fpar_nms
      used_nms_res :: Names
used_nms_res = Names
used_nms0 Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
used_nms' Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Result -> Names
forall a. FreeIn a => a -> Names
freeIn (Body (Aliases rep) -> Result
forall rep. Body rep -> Result
bodyResult Body (Aliases rep)
body)

      -- the last-uses at loop-statement level are the loop free variables that
      -- do not belong to @used_nms0@; this includes the initializers of b), @lu_ini_b@
      lu_arrs :: Names
lu_arrs = Names
used_nms' Names -> Names -> Names
`namesSubtract` Names
used_nms0
  (LUTabFun, Names, Names) -> LastUseM rep (LUTabFun, Names, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
lutab_res, Names
lu_arrs, Names
used_nms_res)
  where
    initHelper :: Names -> (Param dec, SubExp) -> LastUseM rep (Maybe (VName, Names))
initHelper Names
free_and_used (Param dec
fp, SubExp
se) = do
      Names
names <- Names -> LastUseM rep Names
forall rep. Names -> LastUseM rep Names
aliasTransitiveClosure (Names -> LastUseM rep Names) -> Names -> LastUseM rep Names
forall a b. (a -> b) -> a -> b
$ Names -> (VName -> Names) -> Maybe VName -> Names
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Names
forall a. Monoid a => a
mempty VName -> Names
oneName (Maybe VName -> Names) -> Maybe VName -> Names
forall a b. (a -> b) -> a -> b
$ SubExp -> Maybe VName
subExpVar SubExp
se
      if Names
names Names -> Names -> Bool
`namesIntersect` Names
free_and_used
        then Maybe (VName, Names) -> LastUseM rep (Maybe (VName, Names))
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (VName, Names)
forall a. Maybe a
Nothing
        else Maybe (VName, Names) -> LastUseM rep (Maybe (VName, Names))
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (VName, Names) -> LastUseM rep (Maybe (VName, Names)))
-> Maybe (VName, Names) -> LastUseM rep (Maybe (VName, Names))
forall a b. (a -> b) -> a -> b
$ (VName, Names) -> Maybe (VName, Names)
forall a. a -> Maybe a
Just (Ident -> VName
identName (Ident -> VName) -> Ident -> VName
forall a b. (a -> b) -> a -> b
$ Param dec -> Ident
forall dec. Typed dec => Param dec -> Ident
paramIdent Param dec
fp, Names
names)
lastUseExp (Op Op (Aliases rep)
op) Names
used_nms = do
  OpC rep (Aliases rep)
-> Names -> LastUseM rep (LUTabFun, Names, Names)
on_op <- (LastUseReader rep
 -> OpC rep (Aliases rep)
 -> Names
 -> LastUseM rep (LUTabFun, Names, Names))
-> LastUseM
     rep
     (OpC rep (Aliases rep)
      -> Names -> LastUseM rep (LUTabFun, Names, Names))
forall a. (LastUseReader rep -> a) -> LastUseM rep a
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
reader LastUseReader rep
-> OpC rep (Aliases rep)
-> Names
-> LastUseM rep (LUTabFun, Names, Names)
LastUseReader rep -> LastUseOp rep
forall rep. LastUseReader rep -> LastUseOp rep
onOp
  OpC rep (Aliases rep)
-> Names -> LastUseM rep (LUTabFun, Names, Names)
on_op OpC rep (Aliases rep)
Op (Aliases rep)
op Names
used_nms
lastUseExp Exp (Aliases rep)
e Names
used_nms = do
  let free_in_e :: Names
free_in_e = Exp (Aliases rep) -> Names
forall a. FreeIn a => a -> Names
freeIn Exp (Aliases rep)
e
  (Names
used_nms', Names
lu_vars) <- Names -> Names -> LastUseM rep (Names, Names)
forall rep. Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms Names
free_in_e
  (LUTabFun, Names, Names) -> LastUseM rep (LUTabFun, Names, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
forall k a. Map k a
M.empty, Names
lu_vars, Names
used_nms')

lastUseMemOp ::
  (inner (Aliases rep) -> Names -> LastUseM rep (LUTabFun, Names, Names)) ->
  MemOp inner (Aliases rep) ->
  Names ->
  LastUseM rep (LUTabFun, Names, Names)
lastUseMemOp :: forall (inner :: * -> *) rep.
(inner (Aliases rep)
 -> Names -> LastUseM rep (LUTabFun, Names, Names))
-> MemOp inner (Aliases rep)
-> Names
-> LastUseM rep (LUTabFun, Names, Names)
lastUseMemOp inner (Aliases rep)
-> Names -> LastUseM rep (LUTabFun, Names, Names)
_ (Alloc SubExp
se Space
sp) Names
used_nms = do
  let free_in_e :: Names
free_in_e = SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn SubExp
se Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Space -> Names
forall a. FreeIn a => a -> Names
freeIn Space
sp
  (Names
used_nms', Names
lu_vars) <- Names -> Names -> LastUseM rep (Names, Names)
forall rep. Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms Names
free_in_e
  (LUTabFun, Names, Names) -> LastUseM rep (LUTabFun, Names, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
forall k a. Map k a
M.empty, Names
lu_vars, Names
used_nms')
lastUseMemOp inner (Aliases rep)
-> Names -> LastUseM rep (LUTabFun, Names, Names)
onInner (Inner inner (Aliases rep)
op) Names
used_nms = inner (Aliases rep)
-> Names -> LastUseM rep (LUTabFun, Names, Names)
onInner inner (Aliases rep)
op Names
used_nms

lastUseSegOp ::
  (Constraints rep) =>
  SegOp lvl (Aliases rep) ->
  Names ->
  LastUseM rep (LUTabFun, Names, Names)
lastUseSegOp :: forall rep lvl.
Constraints rep =>
SegOp lvl (Aliases rep)
-> Names -> LastUseM rep (LUTabFun, Names, Names)
lastUseSegOp (SegMap lvl
_ SegSpace
_ [Type]
tps KernelBody (Aliases rep)
kbody) Names
used_nms = do
  (Names
used_nms', Names
lu_vars) <- Names -> Names -> LastUseM rep (Names, Names)
forall rep. Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms (Names -> LastUseM rep (Names, Names))
-> Names -> LastUseM rep (Names, Names)
forall a b. (a -> b) -> a -> b
$ [Type] -> Names
forall a. FreeIn a => a -> Names
freeIn [Type]
tps
  (LUTabFun
body_lutab, Names
used_nms'') <- KernelBody (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall rep.
Constraints rep =>
KernelBody (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseKernelBody KernelBody (Aliases rep)
kbody (LUTabFun
forall a. Monoid a => a
mempty, Names
used_nms')
  (LUTabFun, Names, Names) -> LastUseM rep (LUTabFun, Names, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
body_lutab, Names
lu_vars, Names
used_nms' Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
used_nms'')
lastUseSegOp (SegRed lvl
_ SegSpace
_ [Type]
tps KernelBody (Aliases rep)
kbody [SegBinOp (Aliases rep)]
sbos) Names
used_nms = do
  (LUTabFun
lutab_sbo, Names
lu_vars_sbo, Names
used_nms_sbo) <- [SegBinOp (Aliases rep)]
-> Names -> LastUseM rep (LUTabFun, Names, Names)
forall rep.
Constraints rep =>
[SegBinOp (Aliases rep)]
-> Names -> LastUseM rep (LUTabFun, Names, Names)
lastUseSegBinOp [SegBinOp (Aliases rep)]
sbos Names
used_nms
  (Names
used_nms', Names
lu_vars) <- Names -> Names -> LastUseM rep (Names, Names)
forall rep. Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms_sbo (Names -> LastUseM rep (Names, Names))
-> Names -> LastUseM rep (Names, Names)
forall a b. (a -> b) -> a -> b
$ [Type] -> Names
forall a. FreeIn a => a -> Names
freeIn [Type]
tps
  (LUTabFun
body_lutab, Names
used_nms'') <- KernelBody (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall rep.
Constraints rep =>
KernelBody (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseKernelBody KernelBody (Aliases rep)
kbody (LUTabFun
forall a. Monoid a => a
mempty, Names
used_nms')
  (LUTabFun, Names, Names) -> LastUseM rep (LUTabFun, Names, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun -> LUTabFun -> LUTabFun
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union LUTabFun
lutab_sbo LUTabFun
body_lutab, Names
lu_vars Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
lu_vars_sbo, Names
used_nms_sbo Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
used_nms' Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
used_nms'')
lastUseSegOp (SegScan lvl
_ SegSpace
_ [Type]
tps KernelBody (Aliases rep)
kbody [SegBinOp (Aliases rep)]
sbos) Names
used_nms = do
  (LUTabFun
lutab_sbo, Names
lu_vars_sbo, Names
used_nms_sbo) <- [SegBinOp (Aliases rep)]
-> Names -> LastUseM rep (LUTabFun, Names, Names)
forall rep.
Constraints rep =>
[SegBinOp (Aliases rep)]
-> Names -> LastUseM rep (LUTabFun, Names, Names)
lastUseSegBinOp [SegBinOp (Aliases rep)]
sbos Names
used_nms
  (Names
used_nms', Names
lu_vars) <- Names -> Names -> LastUseM rep (Names, Names)
forall rep. Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms_sbo (Names -> LastUseM rep (Names, Names))
-> Names -> LastUseM rep (Names, Names)
forall a b. (a -> b) -> a -> b
$ [Type] -> Names
forall a. FreeIn a => a -> Names
freeIn [Type]
tps
  (LUTabFun
body_lutab, Names
used_nms'') <- KernelBody (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall rep.
Constraints rep =>
KernelBody (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseKernelBody KernelBody (Aliases rep)
kbody (LUTabFun
forall a. Monoid a => a
mempty, Names
used_nms')
  (LUTabFun, Names, Names) -> LastUseM rep (LUTabFun, Names, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun -> LUTabFun -> LUTabFun
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union LUTabFun
lutab_sbo LUTabFun
body_lutab, Names
lu_vars Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
lu_vars_sbo, Names
used_nms_sbo Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
used_nms' Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
used_nms'')
lastUseSegOp (SegHist lvl
_ SegSpace
_ [Type]
tps KernelBody (Aliases rep)
kbody [HistOp (Aliases rep)]
hos) Names
used_nms = do
  (LUTabFun
lutab_sbo, Names
lu_vars_sbo, Names
used_nms_sbo) <- [HistOp (Aliases rep)]
-> Names -> LastUseM rep (LUTabFun, Names, Names)
forall rep.
Constraints rep =>
[HistOp (Aliases rep)]
-> Names -> LastUseM rep (LUTabFun, Names, Names)
lastUseHistOp [HistOp (Aliases rep)]
hos Names
used_nms
  (Names
used_nms', Names
lu_vars) <- Names -> Names -> LastUseM rep (Names, Names)
forall rep. Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms_sbo (Names -> LastUseM rep (Names, Names))
-> Names -> LastUseM rep (Names, Names)
forall a b. (a -> b) -> a -> b
$ [Type] -> Names
forall a. FreeIn a => a -> Names
freeIn [Type]
tps
  (LUTabFun
body_lutab, Names
used_nms'') <- KernelBody (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall rep.
Constraints rep =>
KernelBody (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseKernelBody KernelBody (Aliases rep)
kbody (LUTabFun
forall a. Monoid a => a
mempty, Names
used_nms')
  (LUTabFun, Names, Names) -> LastUseM rep (LUTabFun, Names, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun -> LUTabFun -> LUTabFun
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union LUTabFun
lutab_sbo LUTabFun
body_lutab, Names
lu_vars Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
lu_vars_sbo, Names
used_nms_sbo Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
used_nms' Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
used_nms'')

lastUseGPUOp :: HostOp NoOp (Aliases GPUMem) -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)
lastUseGPUOp :: HostOp NoOp (Aliases GPUMem)
-> Names -> LastUseM GPUMem (LUTabFun, Names, Names)
lastUseGPUOp (GPU.OtherOp NoOp (Aliases GPUMem)
NoOp) Names
used_nms =
  (LUTabFun, Names, Names)
-> LastUseM GPUMem (LUTabFun, Names, Names)
forall a. a -> LastUseM GPUMem a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
forall a. Monoid a => a
mempty, Names
forall a. Monoid a => a
mempty, Names
used_nms)
lastUseGPUOp (SizeOp SizeOp
sop) Names
used_nms = do
  (Names
used_nms', Names
lu_vars) <- Names -> Names -> LastUseM GPUMem (Names, Names)
forall rep. Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms (Names -> LastUseM GPUMem (Names, Names))
-> Names -> LastUseM GPUMem (Names, Names)
forall a b. (a -> b) -> a -> b
$ SizeOp -> Names
forall a. FreeIn a => a -> Names
freeIn SizeOp
sop
  (LUTabFun, Names, Names)
-> LastUseM GPUMem (LUTabFun, Names, Names)
forall a. a -> LastUseM GPUMem a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
forall a. Monoid a => a
mempty, Names
lu_vars, Names
used_nms')
lastUseGPUOp (GPUBody [Type]
tps Body (Aliases GPUMem)
body) Names
used_nms = do
  (Names
used_nms', Names
lu_vars) <- Names -> Names -> LastUseM GPUMem (Names, Names)
forall rep. Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms (Names -> LastUseM GPUMem (Names, Names))
-> Names -> LastUseM GPUMem (Names, Names)
forall a b. (a -> b) -> a -> b
$ [Type] -> Names
forall a. FreeIn a => a -> Names
freeIn [Type]
tps
  (LUTabFun
body_lutab, Names
used_nms'') <- Body (Aliases GPUMem)
-> (LUTabFun, Names) -> LastUseM GPUMem (LUTabFun, Names)
forall rep.
Constraints rep =>
Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseBody Body (Aliases GPUMem)
body (LUTabFun
forall a. Monoid a => a
mempty, Names
used_nms')
  (LUTabFun, Names, Names)
-> LastUseM GPUMem (LUTabFun, Names, Names)
forall a. a -> LastUseM GPUMem a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
body_lutab, Names
lu_vars, Names
used_nms' Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
used_nms'')
lastUseGPUOp (SegOp SegOp SegLevel (Aliases GPUMem)
op) Names
used_nms =
  SegOp SegLevel (Aliases GPUMem)
-> Names -> LastUseM GPUMem (LUTabFun, Names, Names)
forall rep lvl.
Constraints rep =>
SegOp lvl (Aliases rep)
-> Names -> LastUseM rep (LUTabFun, Names, Names)
lastUseSegOp SegOp SegLevel (Aliases GPUMem)
op Names
used_nms

lastUseMCOp :: MCOp NoOp (Aliases MCMem) -> Names -> LastUseM MCMem (LUTabFun, Names, Names)
lastUseMCOp :: MCOp NoOp (Aliases MCMem)
-> Names -> LastUseM MCMem (LUTabFun, Names, Names)
lastUseMCOp (MC.OtherOp NoOp (Aliases MCMem)
NoOp) Names
used_nms =
  (LUTabFun, Names, Names) -> LastUseM MCMem (LUTabFun, Names, Names)
forall a. a -> LastUseM MCMem a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
forall a. Monoid a => a
mempty, Names
forall a. Monoid a => a
mempty, Names
used_nms)
lastUseMCOp (MC.ParOp Maybe (SegOp () (Aliases MCMem))
par_op SegOp () (Aliases MCMem)
op) Names
used_nms = do
  (LUTabFun
lutab_par_op, Names
lu_vars_par_op, Names
used_names_par_op) <-
    LastUseM MCMem (LUTabFun, Names, Names)
-> (SegOp () (Aliases MCMem)
    -> LastUseM MCMem (LUTabFun, Names, Names))
-> Maybe (SegOp () (Aliases MCMem))
-> LastUseM MCMem (LUTabFun, Names, Names)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ((LUTabFun, Names, Names) -> LastUseM MCMem (LUTabFun, Names, Names)
forall a. a -> LastUseM MCMem a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun, Names, Names)
forall a. Monoid a => a
mempty) (SegOp () (Aliases MCMem)
-> Names -> LastUseM MCMem (LUTabFun, Names, Names)
forall rep lvl.
Constraints rep =>
SegOp lvl (Aliases rep)
-> Names -> LastUseM rep (LUTabFun, Names, Names)
`lastUseSegOp` Names
used_nms) Maybe (SegOp () (Aliases MCMem))
par_op
  (LUTabFun
lutab_op, Names
lu_vars_op, Names
used_names_op) <-
    SegOp () (Aliases MCMem)
-> Names -> LastUseM MCMem (LUTabFun, Names, Names)
forall rep lvl.
Constraints rep =>
SegOp lvl (Aliases rep)
-> Names -> LastUseM rep (LUTabFun, Names, Names)
lastUseSegOp SegOp () (Aliases MCMem)
op Names
used_nms
  (LUTabFun, Names, Names) -> LastUseM MCMem (LUTabFun, Names, Names)
forall a. a -> LastUseM MCMem a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    ( LUTabFun
lutab_par_op LUTabFun -> LUTabFun -> LUTabFun
forall a. Semigroup a => a -> a -> a
<> LUTabFun
lutab_op,
      Names
lu_vars_par_op Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
lu_vars_op,
      Names
used_names_par_op Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
used_names_op
    )

lastUseSegBinOp ::
  (Constraints rep) =>
  [SegBinOp (Aliases rep)] ->
  Names ->
  LastUseM rep (LUTabFun, Names, Names)
lastUseSegBinOp :: forall rep.
Constraints rep =>
[SegBinOp (Aliases rep)]
-> Names -> LastUseM rep (LUTabFun, Names, Names)
lastUseSegBinOp [SegBinOp (Aliases rep)]
sbos Names
used_nms = do
  ([LUTabFun]
lutab, [Names]
lu_vars, [Names]
used_nms') <- [(LUTabFun, Names, Names)] -> ([LUTabFun], [Names], [Names])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(LUTabFun, Names, Names)] -> ([LUTabFun], [Names], [Names]))
-> LastUseM rep [(LUTabFun, Names, Names)]
-> LastUseM rep ([LUTabFun], [Names], [Names])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SegBinOp (Aliases rep) -> LastUseM rep (LUTabFun, Names, Names))
-> [SegBinOp (Aliases rep)]
-> LastUseM rep [(LUTabFun, Names, Names)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SegBinOp (Aliases rep) -> LastUseM rep (LUTabFun, Names, Names)
helper [SegBinOp (Aliases rep)]
sbos
  (LUTabFun, Names, Names) -> LastUseM rep (LUTabFun, Names, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([LUTabFun] -> LUTabFun
forall a. Monoid a => [a] -> a
mconcat [LUTabFun]
lutab, [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat [Names]
lu_vars, [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat [Names]
used_nms')
  where
    helper :: SegBinOp (Aliases rep) -> LastUseM rep (LUTabFun, Names, Names)
helper (SegBinOp Commutativity
_ l :: Lambda (Aliases rep)
l@(Lambda [LParam (Aliases rep)]
_ [Type]
_ Body (Aliases rep)
body) [SubExp]
neutral ShapeBase SubExp
shp) = Lambda (Aliases rep)
-> LastUseM rep (LUTabFun, Names, Names)
-> LastUseM rep (LUTabFun, Names, Names)
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Lambda (Aliases rep)
l (LastUseM rep (LUTabFun, Names, Names)
 -> LastUseM rep (LUTabFun, Names, Names))
-> LastUseM rep (LUTabFun, Names, Names)
-> LastUseM rep (LUTabFun, Names, Names)
forall a b. (a -> b) -> a -> b
$ do
      (Names
used_nms', Names
lu_vars) <- Names -> Names -> LastUseM rep (Names, Names)
forall rep. Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms (Names -> LastUseM rep (Names, Names))
-> Names -> LastUseM rep (Names, Names)
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Names
forall a. FreeIn a => a -> Names
freeIn [SubExp]
neutral Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> ShapeBase SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn ShapeBase SubExp
shp
      (LUTabFun
body_lutab, Names
used_nms'') <- Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall rep.
Constraints rep =>
Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseBody Body (Aliases rep)
body (LUTabFun
forall a. Monoid a => a
mempty, Names
used_nms')
      (LUTabFun, Names, Names) -> LastUseM rep (LUTabFun, Names, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
body_lutab, Names
lu_vars, Names
used_nms'')

lastUseHistOp ::
  (Constraints rep) =>
  [HistOp (Aliases rep)] ->
  Names ->
  LastUseM rep (LUTabFun, Names, Names)
lastUseHistOp :: forall rep.
Constraints rep =>
[HistOp (Aliases rep)]
-> Names -> LastUseM rep (LUTabFun, Names, Names)
lastUseHistOp [HistOp (Aliases rep)]
hos Names
used_nms = do
  ([LUTabFun]
lutab, [Names]
lu_vars, [Names]
used_nms') <- [(LUTabFun, Names, Names)] -> ([LUTabFun], [Names], [Names])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(LUTabFun, Names, Names)] -> ([LUTabFun], [Names], [Names]))
-> LastUseM rep [(LUTabFun, Names, Names)]
-> LastUseM rep ([LUTabFun], [Names], [Names])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (HistOp (Aliases rep) -> LastUseM rep (LUTabFun, Names, Names))
-> [HistOp (Aliases rep)]
-> LastUseM rep [(LUTabFun, Names, Names)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM HistOp (Aliases rep) -> LastUseM rep (LUTabFun, Names, Names)
helper [HistOp (Aliases rep)]
hos
  (LUTabFun, Names, Names) -> LastUseM rep (LUTabFun, Names, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([LUTabFun] -> LUTabFun
forall a. Monoid a => [a] -> a
mconcat [LUTabFun]
lutab, [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat [Names]
lu_vars, [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat [Names]
used_nms')
  where
    helper :: HistOp (Aliases rep) -> LastUseM rep (LUTabFun, Names, Names)
helper (HistOp ShapeBase SubExp
shp SubExp
rf [VName]
dest [SubExp]
neutral ShapeBase SubExp
shp' l :: Lambda (Aliases rep)
l@(Lambda [LParam (Aliases rep)]
_ [Type]
_ Body (Aliases rep)
body)) = Lambda (Aliases rep)
-> LastUseM rep (LUTabFun, Names, Names)
-> LastUseM rep (LUTabFun, Names, Names)
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Lambda (Aliases rep)
l (LastUseM rep (LUTabFun, Names, Names)
 -> LastUseM rep (LUTabFun, Names, Names))
-> LastUseM rep (LUTabFun, Names, Names)
-> LastUseM rep (LUTabFun, Names, Names)
forall a b. (a -> b) -> a -> b
$ do
      (Names
used_nms', Names
lu_vars) <- Names -> Names -> LastUseM rep (Names, Names)
forall rep. Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms (Names -> LastUseM rep (Names, Names))
-> Names -> LastUseM rep (Names, Names)
forall a b. (a -> b) -> a -> b
$ ShapeBase SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn ShapeBase SubExp
shp Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn SubExp
rf Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [VName] -> Names
forall a. FreeIn a => a -> Names
freeIn [VName]
dest Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> Names
forall a. FreeIn a => a -> Names
freeIn [SubExp]
neutral Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> ShapeBase SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn ShapeBase SubExp
shp'
      (LUTabFun
body_lutab, Names
used_nms'') <- Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
forall rep.
Constraints rep =>
Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseBody Body (Aliases rep)
body (LUTabFun
forall a. Monoid a => a
mempty, Names
used_nms')
      (LUTabFun, Names, Names) -> LastUseM rep (LUTabFun, Names, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
body_lutab, Names
lu_vars, Names
used_nms'')

lastUseSeqOp :: Op (Aliases SeqMem) -> Names -> LastUseM SeqMem (LUTabFun, Names, Names)
lastUseSeqOp :: LastUseOp SeqMem
lastUseSeqOp (Alloc SubExp
se Space
sp) Names
used_nms = do
  let free_in_e :: Names
free_in_e = SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn SubExp
se Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Space -> Names
forall a. FreeIn a => a -> Names
freeIn Space
sp
  (Names
used_nms', Names
lu_vars) <- Names -> Names -> LastUseM SeqMem (Names, Names)
forall rep. Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms Names
free_in_e
  (LUTabFun, Names, Names)
-> LastUseM SeqMem (LUTabFun, Names, Names)
forall a. a -> LastUseM SeqMem a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
forall a. Monoid a => a
mempty, Names
lu_vars, Names
used_nms')
lastUseSeqOp (Inner NoOp (Aliases SeqMem)
NoOp) Names
used_nms = do
  (LUTabFun, Names, Names)
-> LastUseM SeqMem (LUTabFun, Names, Names)
forall a. a -> LastUseM SeqMem a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
forall a. Monoid a => a
mempty, Names
forall a. Monoid a => a
mempty, Names
used_nms)

------------------------------------------------------

-- | Given already used names and newly encountered 'Names', return an updated
-- set used names and the set of names that were last used here.
--
-- For a given name @x@ in the new uses, if neither @x@ nor any of its aliases
-- are present in the set of used names, this is a last use of @x@.
lastUsedInNames ::
  -- | Used names
  Names ->
  -- | New uses
  Names ->
  LastUseM rep (Names, Names)
lastUsedInNames :: forall rep. Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms Names
new_uses = do
  -- a use of an argument x is also a use of any variable in x alias set
  -- so we update the alias-based transitive-closure of used names.
  Names
new_uses_with_aliases <- Names -> LastUseM rep Names
forall rep. Names -> LastUseM rep Names
aliasTransitiveClosure Names
new_uses
  -- if neither a variable x, nor any of its alias set have been used before (in
  -- the backward traversal), then it is a last use of both that variable and
  -- all other variables in its alias set
  [VName]
last_uses <- (VName -> LastUseM rep Bool) -> [VName] -> LastUseM rep [VName]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM VName -> LastUseM rep Bool
isLastUse ([VName] -> LastUseM rep [VName])
-> [VName] -> LastUseM rep [VName]
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList Names
new_uses
  Names
last_uses' <- Names -> LastUseM rep Names
forall rep. Names -> LastUseM rep Names
aliasTransitiveClosure (Names -> LastUseM rep Names) -> Names -> LastUseM rep Names
forall a b. (a -> b) -> a -> b
$ [VName] -> Names
namesFromList [VName]
last_uses
  (Names, Names) -> LastUseM rep (Names, Names)
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Names
used_nms Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
new_uses_with_aliases, Names
last_uses')
  where
    isLastUse :: VName -> LastUseM rep Bool
isLastUse VName
x = do
      Names
with_aliases <- Names -> LastUseM rep Names
forall rep. Names -> LastUseM rep Names
aliasTransitiveClosure (Names -> LastUseM rep Names) -> Names -> LastUseM rep Names
forall a b. (a -> b) -> a -> b
$ VName -> Names
oneName VName
x
      Bool -> LastUseM rep Bool
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> LastUseM rep Bool) -> Bool -> LastUseM rep Bool
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Names
with_aliases Names -> Names -> Bool
`namesIntersect` Names
used_nms

-- | Compute the transitive closure of the aliases of a set of 'Names'.
aliasTransitiveClosure :: Names -> LastUseM rep Names
aliasTransitiveClosure :: forall rep. Names -> LastUseM rep Names
aliasTransitiveClosure Names
args = do
  Names
res <- (Names -> Names -> Names) -> Names -> [Names] -> Names
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
(<>) Names
args ([Names] -> Names) -> LastUseM rep [Names] -> LastUseM rep Names
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> LastUseM rep Names) -> [VName] -> LastUseM rep [Names]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> LastUseM rep Names
forall rep. VName -> LastUseM rep Names
aliasLookup (Names -> [VName]
namesToList Names
args)
  if Names
res Names -> Names -> Bool
forall a. Eq a => a -> a -> Bool
== Names
args
    then Names -> LastUseM rep Names
forall a. a -> LastUseM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Names
res
    else Names -> LastUseM rep Names
forall rep. Names -> LastUseM rep Names
aliasTransitiveClosure Names
res

-- | For each 'PatElem' in the 'Pat', add its aliases to the 'AliasTab' in
-- 'LastUseM'. Additionally, 'Names' are added as aliases of all the 'PatElemT'.
updateAliasing ::
  (AliasesOf dec) =>
  -- | Extra names that all 'PatElem' should alias.
  Names ->
  -- | Pattern to process
  Pat dec ->
  LastUseM rep ()
updateAliasing :: forall dec rep.
AliasesOf dec =>
Names -> Pat dec -> LastUseM rep ()
updateAliasing Names
extra_aliases =
  (PatElem dec -> LastUseM rep ())
-> [PatElem dec] -> LastUseM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ PatElem dec -> LastUseM rep ()
forall dec rep. AliasesOf dec => PatElem dec -> LastUseM rep ()
update ([PatElem dec] -> LastUseM rep ())
-> (Pat dec -> [PatElem dec]) -> Pat dec -> LastUseM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat dec -> [PatElem dec]
forall dec. Pat dec -> [PatElem dec]
patElems
  where
    update :: (AliasesOf dec) => PatElem dec -> LastUseM rep ()
    update :: forall dec rep. AliasesOf dec => PatElem dec -> LastUseM rep ()
update (PatElem VName
name dec
dec) = do
      let aliases :: Names
aliases = dec -> Names
forall a. AliasesOf a => a -> Names
aliasesOf dec
dec
      Names
aliases' <- Names -> LastUseM rep Names
forall rep. Names -> LastUseM rep Names
aliasTransitiveClosure (Names -> LastUseM rep Names) -> Names -> LastUseM rep Names
forall a b. (a -> b) -> a -> b
$ Names
extra_aliases Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
aliases
      (LUTabFun -> LUTabFun) -> LastUseM rep ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((LUTabFun -> LUTabFun) -> LastUseM rep ())
-> (LUTabFun -> LUTabFun) -> LastUseM rep ()
forall a b. (a -> b) -> a -> b
$ VName -> Names -> LUTabFun -> LUTabFun
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
name Names
aliases'