{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- | 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.
--
-- This pass is different from "Futhark.Analysis.LastUse" in that memory blocks
-- are used to alias arrays. For instance, an 'Update' will not result in a last
-- use of the array being updated, because the result lives in the same memory.
module Futhark.Optimise.ArrayShortCircuiting.LastUse
  ( lastUseSeqMem,
    lastUseGPUMem,
    lastUseMCMem,
    LUTabFun,
    LUTabProg,
  )
where

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 {k} (rep :: k). LastUseReader rep -> LastUseOp rep
onOp :: LastUseOp rep,
    forall {k} (rep :: k). 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
    ( forall a. a -> LastUseM rep a
forall k (rep :: k). Applicative (LastUseM rep)
forall k (rep :: k) a. a -> LastUseM rep a
forall k (rep :: k) a b.
LastUseM rep a -> LastUseM rep b -> LastUseM rep b
forall k (rep :: k) 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 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
return :: forall a. a -> LastUseM rep a
$creturn :: forall k (rep :: k) a. a -> LastUseM rep a
>> :: forall a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep b
$c>> :: forall k (rep :: k) a b.
LastUseM rep a -> LastUseM rep b -> LastUseM rep b
>>= :: forall a b.
LastUseM rep a -> (a -> LastUseM rep b) -> LastUseM rep b
$c>>= :: forall k (rep :: k) a b.
LastUseM rep a -> (a -> LastUseM rep b) -> LastUseM rep b
Monad,
      forall k (rep :: k) a b. a -> LastUseM rep b -> LastUseM rep a
forall k (rep :: k) a b.
(a -> b) -> LastUseM rep a -> LastUseM rep b
forall a b. a -> LastUseM rep b -> LastUseM rep a
forall 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
<$ :: forall a b. a -> LastUseM rep b -> LastUseM rep a
$c<$ :: forall k (rep :: k) a b. a -> LastUseM rep b -> LastUseM rep a
fmap :: forall a b. (a -> b) -> LastUseM rep a -> LastUseM rep b
$cfmap :: forall k (rep :: k) a b.
(a -> b) -> LastUseM rep a -> LastUseM rep b
Functor,
      forall a. a -> LastUseM rep a
forall k (rep :: k). Functor (LastUseM rep)
forall k (rep :: k) a. a -> LastUseM rep a
forall k (rep :: k) a b.
LastUseM rep a -> LastUseM rep b -> LastUseM rep a
forall k (rep :: k) a b.
LastUseM rep a -> LastUseM rep b -> LastUseM rep b
forall k (rep :: k) a b.
LastUseM rep (a -> b) -> LastUseM rep a -> LastUseM rep b
forall k (rep :: k) 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 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 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
<* :: forall a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep a
$c<* :: forall k (rep :: k) a b.
LastUseM rep a -> LastUseM rep b -> LastUseM rep a
*> :: forall a b. LastUseM rep a -> LastUseM rep b -> LastUseM rep b
$c*> :: forall k (rep :: k) a b.
LastUseM rep a -> LastUseM rep b -> LastUseM rep b
liftA2 :: forall a b c.
(a -> b -> c) -> LastUseM rep a -> LastUseM rep b -> LastUseM rep c
$cliftA2 :: forall k (rep :: k) a b c.
(a -> b -> c) -> LastUseM rep a -> LastUseM rep b -> LastUseM rep c
<*> :: forall a b.
LastUseM rep (a -> b) -> LastUseM rep a -> LastUseM rep b
$c<*> :: forall k (rep :: k) a b.
LastUseM rep (a -> b) -> LastUseM rep a -> LastUseM rep b
pure :: forall a. a -> LastUseM rep a
$cpure :: forall k (rep :: k) a. a -> LastUseM rep a
Applicative,
      MonadReader (LastUseReader rep),
      MonadState AliasTab
    )

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

instance
  (RepTypes rep, CanBeAliased (Op 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) = forall {k} (rep :: k) a.
StateT LUTabFun (Reader (LastUseReader rep)) a -> LastUseM rep a
LastUseM forall a b. (a -> b) -> a -> b
$ do
    forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (\LastUseReader rep
rd -> LastUseReader rep
rd {scope :: Scope (Aliases rep)
scope = forall {k} (rep :: k). LastUseReader rep -> Scope (Aliases rep)
scope LastUseReader rep
rd forall a. Semigroup a => a -> a -> a
<> Scope (Aliases rep)
sc}) StateT LUTabFun (Reader (LastUseReader rep)) a
m

type Constraints rep =
  ( LocalScope (Aliases rep) (LastUseM rep),
    ASTRep rep,
    FreeIn (OpWithAliases (Op rep)),
    HasMemBlock (Aliases rep),
    CanBeAliased (Op rep)
  )

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

aliasLookup :: VName -> LastUseM rep Names
aliasLookup :: forall {k} (rep :: k). VName -> LastUseM rep Names
aliasLookup VName
vname =
  forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a -> a
fromMaybe forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 {k} (rep :: k).
Constraints rep =>
Prog (Aliases rep) -> LastUseM rep LUTabProg
lastUseProg Prog (Aliases rep)
prog =
  let bound_in_consts :: Names
bound_in_consts =
        forall {k} (rep :: k). Prog rep -> Stms rep
progConsts Prog (Aliases rep)
prog
          forall a b. a -> (a -> b) -> b
& forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall dec. Pat dec -> [VName]
patNames forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). Stm rep -> Pat (LetDec rep)
stmPat)
          forall a b. a -> (a -> b) -> b
& [VName] -> Names
namesFromList
      consts :: Stms (Aliases rep)
consts = forall {k} (rep :: k). Prog rep -> Stms rep
progConsts Prog (Aliases rep)
prog
      funs :: [FunDef (Aliases rep)]
funs = forall {k} (rep :: k). Prog rep -> [FunDef rep]
progFuns Prog (Aliases rep)
prog
   in forall {k} (rep :: k) a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Stms (Aliases rep)
consts forall a b. (a -> b) -> a -> b
$ do
        (LUTabFun
consts_lu, Names
_) <- forall {k} (rep :: k).
Constraints rep =>
Stms (Aliases rep)
-> (LUTabFun, Names) -> [VName] -> LastUseM rep (LUTabFun, Names)
lastUseStms Stms (Aliases rep)
consts forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty
        [LUTabFun]
lus <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} (rep :: k).
Constraints rep =>
Names -> FunDef (Aliases rep) -> LastUseM rep LUTabFun
lastUseFun Names
bound_in_consts) [FunDef (Aliases rep)]
funs
        forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
consts_lu, 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 {k} (rep :: k). FunDef rep -> Name
funDefName [FunDef (Aliases rep)]
funs) [LUTabFun]
lus)

lastUseFun ::
  Constraints rep =>
  Names ->
  FunDef (Aliases rep) ->
  LastUseM rep LUTabFun
lastUseFun :: forall {k} (rep :: k).
Constraints rep =>
Names -> FunDef (Aliases rep) -> LastUseM rep LUTabFun
lastUseFun Names
bound_in_consts FunDef (Aliases rep)
f =
  forall {k} (rep :: k) a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf FunDef (Aliases rep)
f forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (rep :: k).
Constraints rep =>
Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseBody (forall {k} (rep :: k). FunDef rep -> Body rep
funDefBody FunDef (Aliases rep)
f) (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 = forall {k} (rep :: k) a. LastUseOp rep -> LastUseM rep a -> a
runLastUseM Op (Aliases SeqMem)
-> Names -> LastUseM SeqMem (LUTabFun, Names, Names)
lastUseSeqOp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k).
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 = forall {k} (rep :: k) a. LastUseOp rep -> LastUseM rep a -> a
runLastUseM (forall {k} inner (rep :: k).
(inner -> Names -> LastUseM rep (LUTabFun, Names, Names))
-> MemOp inner -> Names -> LastUseM rep (LUTabFun, Names, Names)
lastUseMemOp HostOp (Aliases GPUMem) ()
-> Names -> LastUseM GPUMem (LUTabFun, Names, Names)
lastUseGPUOp) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k).
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 = forall {k} (rep :: k) a. LastUseOp rep -> LastUseM rep a -> a
runLastUseM (forall {k} inner (rep :: k).
(inner -> Names -> LastUseM rep (LUTabFun, Names, Names))
-> MemOp inner -> Names -> LastUseM rep (LUTabFun, Names, Names)
lastUseMemOp MCOp (Aliases MCMem) ()
-> Names -> LastUseM MCMem (LUTabFun, Names, Names)
lastUseMCOp) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k).
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 {k} (rep :: k).
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.
  forall {k} (rep :: k) a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Stms (Aliases rep)
stms forall a b. (a -> b) -> a -> b
$ do
    (LUTabFun
lutab', Names
_) <-
      forall {k} (rep :: k).
Constraints rep =>
Stms (Aliases rep)
-> (LUTabFun, Names) -> [VName] -> LastUseM rep (LUTabFun, Names)
lastUseStms Stms (Aliases rep)
stms (LUTabFun
lutab, Names
used_nms) forall a b. (a -> b) -> a -> b
$
        Names -> [VName]
namesToList forall a b. (a -> b) -> a -> b
$
          forall a. FreeIn a => a -> Names
freeIn forall a b. (a -> b) -> a -> b
$
            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 <- forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn Body (Aliases rep)
bdy
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
lutab', Names
used_nms 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 {k} (rep :: k).
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) =
  forall {k} (rep :: k) a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Stms (Aliases rep)
stms 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
_) <-
      forall {k} (rep :: k).
Constraints rep =>
Stms (Aliases rep)
-> (LUTabFun, Names) -> [VName] -> LastUseM rep (LUTabFun, Names)
lastUseStms Stms (Aliases rep)
stms (LUTabFun
lutab, Names
used_nms) forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList forall a b. (a -> b) -> a -> b
$ 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 <- forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn KernelBody (Aliases rep)
bdy
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
lutab', Names
used_nms 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 {k} (rep :: k).
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 <- forall (m :: * -> *) b a.
(Monad m, Monoid b) =>
(a -> m b) -> [a] -> m b
concatMapM forall {k} (rep :: k). VName -> LastUseM rep Names
aliasLookup [VName]
res_nms
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
lutab, Names
nms forall a. Semigroup a => a -> a -> a
<> Names
aliases 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 =
  forall {k} (rep :: k) a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Stm (Aliases rep)
stm 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)
_ -> forall a. Monoid a => a
mempty
    -- We build up aliases top-down
    forall {k} dec (rep :: k).
AliasesOf dec =>
Names -> Pat dec -> LastUseM rep ()
updateAliasing Names
extra_alias Pat (LetDec (Aliases rep))
pat
    -- But compute last use bottom-up
    (LUTabFun
lutab', Names
nms') <- forall {k} (rep :: k).
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'') <- forall {k} (rep :: k).
Constraints rep =>
Stm (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseStm Stm (Aliases rep)
stm (LUTabFun
lutab', Names
nms')
    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 {k} (rep :: k).
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') <- forall {k} (rep :: k).
Constraints rep =>
Exp (Aliases rep) -> Names -> LastUseM rep (LUTabFun, Names, Names)
lastUseExp Exp (Aliases rep)
e Names
used_nms
  Scope (Aliases rep)
sc <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall {k} (rep :: k). LastUseReader rep -> Scope (Aliases rep)
scope
  let lu_mems :: Names
lu_mems =
        Names -> [VName]
namesToList Names
last_uses
          forall a b. a -> (a -> b) -> b
& forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (forall {k} (rep :: k).
HasMemBlock rep =>
VName -> Scope rep -> Maybe ArrayMemBound
`getScopeMemInfo` Scope (Aliases rep)
sc)
          forall a b. a -> (a -> b) -> b
& forall a b. (a -> b) -> [a] -> [b]
map ArrayMemBound -> VName
memName
          forall a b. a -> (a -> b) -> b
& [VName] -> Names
namesFromList
          forall a b. a -> (a -> b) -> b
& 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 = forall dec. Pat dec -> [VName]
patNames 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'' =
        forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union LUTabFun
lutab' forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (forall a. [a] -> a
head [VName]
patnms) (Names
last_uses forall a. Semigroup a => a -> a -> a
<> Names
lu_mems) LUTabFun
lutab
  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 {k} (rep :: k).
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) <-
    forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap forall a. Monoid a => [a] -> a
mconcat forall a. Monoid a => [a] -> a
mconcat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [(a, b)] -> ([a], [b])
unzip
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall {k} (rep :: k).
Constraints rep =>
Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseBody (forall k a. Map k a
M.empty, Names
used_nms) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall body. Case body -> body
caseBody) [Case (Body (Aliases rep))]
cases
  (LUTabFun
lutab', Names
body_used_nms) <- forall {k} (rep :: k).
Constraints rep =>
Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseBody Body (Aliases rep)
body (forall k a. Map k a
M.empty, Names
used_nms)
  let free_in_body :: Names
free_in_body = forall a. FreeIn a => a -> Names
freeIn Body (Aliases rep)
body
  let free_in_cases :: Names
free_in_cases = forall a. FreeIn a => a -> Names
freeIn [Case (Body (Aliases rep))]
cases
  let used_nms' :: Names
used_nms' = Names
used_cases forall a. Semigroup a => a -> a -> a
<> Names
body_used_nms
  (Names
_, Names
last_used_arrs) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms forall a b. (a -> b) -> a -> b
$ Names
free_in_body forall a. Semigroup a => a -> a -> a
<> Names
free_in_cases
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (LUTabFun
lutab_cases forall a. Semigroup a => a -> a -> a
<> LUTabFun
lutab', Names
last_used_arrs, Names
used_nms')
lastUseExp (DoLoop [(FParam (Aliases rep), SubExp)]
var_ses LoopForm (Aliases rep)
lf Body (Aliases rep)
body) Names
used_nms0 = forall {k} (rep :: k) a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf LoopForm (Aliases rep)
lf forall a b. (a -> b) -> a -> b
$ do
  Names
free_in_body <- forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure forall a b. (a -> b) -> a -> b
$ 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 <- forall a. [Maybe a] -> [a]
catMaybes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {dec} {rep :: k}.
Typed dec =>
Names -> (Param dec, SubExp) -> LastUseM rep (Maybe (VName, Names))
initHelper (Names
free_in_body forall a. Semigroup a => a -> a -> a
<> Names
used_nms0)) [(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 (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(VName, Names)]
var_inis)
      used_nms :: Names
used_nms = Names
used_nms0 forall a. Semigroup a => a -> a -> a
<> Names
free_in_body' forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn (forall {k} (rep :: k). Body rep -> Result
bodyResult Body (Aliases rep)
body)
  (LUTabFun
body_lutab, Names
_) <- forall {k} (rep :: k).
Constraints rep =>
Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseBody Body (Aliases rep)
body (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 forall a. Semigroup a => a -> a -> a
<> 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 forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (Ident -> VName
identName forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. Typed dec => Param dec -> Ident
paramIdent forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(FParam (Aliases rep), SubExp)]
var_ses
      used_nms' :: Names
used_nms' = (Names
free_in_body forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(FParam (Aliases rep), SubExp)]
var_ses)) Names -> Names -> Names
`namesSubtract` Names
fpar_nms
      used_nms_res :: Names
used_nms_res = Names
used_nms0 forall a. Semigroup a => a -> a -> a
<> Names
used_nms' forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn (forall {k} (rep :: k). 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
  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 <- forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure forall a b. (a -> b) -> a -> b
$ forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty VName -> Names
oneName 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 forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
        else forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (Ident -> VName
identName forall a b. (a -> b) -> a -> b
$ forall dec. Typed dec => Param dec -> Ident
paramIdent Param dec
fp, Names
names)
lastUseExp (Op Op (Aliases rep)
op) Names
used_nms = do
  OpWithAliases (Op rep)
-> Names -> LastUseM rep (LUTabFun, Names, Names)
on_op <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
reader forall {k} (rep :: k). LastUseReader rep -> LastUseOp rep
onOp
  OpWithAliases (Op rep)
-> Names -> LastUseM rep (LUTabFun, Names, Names)
on_op Op (Aliases rep)
op Names
used_nms
lastUseExp Exp (Aliases rep)
e Names
used_nms = do
  let free_in_e :: Names
free_in_e = forall a. FreeIn a => a -> Names
freeIn Exp (Aliases rep)
e
  (Names
used_nms', Names
lu_vars) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms Names
free_in_e
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall k a. Map k a
M.empty, Names
lu_vars, Names
used_nms')

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

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

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

lastUseMCOp :: MCOp (Aliases MCMem) () -> Names -> LastUseM MCMem (LUTabFun, Names, Names)
lastUseMCOp :: MCOp (Aliases MCMem) ()
-> Names -> LastUseM MCMem (LUTabFun, Names, Names)
lastUseMCOp (MC.OtherOp ()) Names
used_nms =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Monoid a => a
mempty, 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) <-
    forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty) (forall {k} (rep :: k) 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) <-
    forall {k} (rep :: k) lvl.
Constraints rep =>
SegOp lvl (Aliases rep)
-> Names -> LastUseM rep (LUTabFun, Names, Names)
lastUseSegOp SegOp () (Aliases MCMem)
op Names
used_nms
  forall (f :: * -> *) a. Applicative f => a -> f a
pure
    ( LUTabFun
lutab_par_op forall a. Semigroup a => a -> a -> a
<> LUTabFun
lutab_op,
      Names
lu_vars_par_op forall a. Semigroup a => a -> a -> a
<> Names
lu_vars_op,
      Names
used_names_par_op 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 {k} (rep :: k).
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') <- forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SegBinOp (Aliases rep) -> LastUseM rep (LUTabFun, Names, Names)
helper [SegBinOp (Aliases rep)]
sbos
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Monoid a => [a] -> a
mconcat [LUTabFun]
lutab, forall a. Monoid a => [a] -> a
mconcat [Names]
lu_vars, 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)]
_ Body (Aliases rep)
body [Type]
_) [SubExp]
neutral ShapeBase SubExp
shp) = forall {k} (rep :: k) a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Lambda (Aliases rep)
l forall a b. (a -> b) -> a -> b
$ do
      (Names
used_nms', Names
lu_vars) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn [SubExp]
neutral forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn ShapeBase SubExp
shp
      (LUTabFun
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
Constraints rep =>
Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseBody Body (Aliases rep)
body (forall a. Monoid a => a
mempty, Names
used_nms')
      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 {k} (rep :: k).
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') <- forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM HistOp (Aliases rep) -> LastUseM rep (LUTabFun, Names, Names)
helper [HistOp (Aliases rep)]
hos
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Monoid a => [a] -> a
mconcat [LUTabFun]
lutab, forall a. Monoid a => [a] -> a
mconcat [Names]
lu_vars, 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)]
_ Body (Aliases rep)
body [Type]
_)) = forall {k} (rep :: k) a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Lambda (Aliases rep)
l forall a b. (a -> b) -> a -> b
$ do
      (Names
used_nms', Names
lu_vars) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn ShapeBase SubExp
shp forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn SubExp
rf forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn [VName]
dest forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn [SubExp]
neutral forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn ShapeBase SubExp
shp'
      (LUTabFun
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
Constraints rep =>
Body (Aliases rep)
-> (LUTabFun, Names) -> LastUseM rep (LUTabFun, Names)
lastUseBody Body (Aliases rep)
body (forall a. Monoid a => a
mempty, Names
used_nms')
      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 :: Op (Aliases SeqMem)
-> Names -> LastUseM SeqMem (LUTabFun, Names, Names)
lastUseSeqOp (Alloc SubExp
se Space
sp) Names
used_nms = do
  let free_in_e :: Names
free_in_e = forall a. FreeIn a => a -> Names
freeIn SubExp
se forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn Space
sp
  (Names
used_nms', Names
lu_vars) <- forall {k} (rep :: k).
Names -> Names -> LastUseM rep (Names, Names)
lastUsedInNames Names
used_nms Names
free_in_e
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Monoid a => a
mempty, Names
lu_vars, Names
used_nms')
lastUseSeqOp (Inner ()) Names
used_nms = do
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Monoid a => a
mempty, 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 {k} (rep :: k).
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 <- forall {k} (rep :: k). 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 <- forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM VName -> LastUseM rep Bool
isLastUse forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList Names
new_uses
  Names
last_uses' <- forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure forall a b. (a -> b) -> a -> b
$ [VName] -> Names
namesFromList [VName]
last_uses
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (Names
used_nms 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 <- forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure forall a b. (a -> b) -> a -> b
$ VName -> Names
oneName VName
x
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not 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 {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure Names
args = do
  Names
res <- forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall a. Semigroup a => a -> a -> a
(<>) Names
args forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {k} (rep :: k). VName -> LastUseM rep Names
aliasLookup (Names -> [VName]
namesToList Names
args)
  if Names
res forall a. Eq a => a -> a -> Bool
== Names
args
    then forall (f :: * -> *) a. Applicative f => a -> f a
pure Names
res
    else forall {k} (rep :: k). 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 {k} dec (rep :: k).
AliasesOf dec =>
Names -> Pat dec -> LastUseM rep ()
updateAliasing Names
extra_aliases =
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ forall {k} dec (rep :: k).
AliasesOf dec =>
PatElem dec -> LastUseM rep ()
update forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. Pat dec -> [PatElem dec]
patElems
  where
    update :: AliasesOf dec => PatElem dec -> LastUseM rep ()
    update :: forall {k} dec (rep :: k).
AliasesOf dec =>
PatElem dec -> LastUseM rep ()
update (PatElem VName
name dec
dec) = do
      let aliases :: Names
aliases = forall a. AliasesOf a => a -> Names
aliasesOf dec
dec
      Names
aliases' <- forall {k} (rep :: k). Names -> LastUseM rep Names
aliasTransitiveClosure forall a b. (a -> b) -> a -> b
$ Names
extra_aliases forall a. Semigroup a => a -> a -> a
<> Names
aliases
      forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
name Names
aliases'