{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# 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.
--
-- 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, lastUsePrg, lastUsePrgGPU, lastUseGPUMem) where

import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.Bifunctor (bimap)
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.SeqMem
import Futhark.Optimise.ArrayShortCircuiting.DataStructs
import Futhark.Util

-- | 'LastUseReader' allows us to abstract over representations by supplying the
-- 'onOp' function.
newtype LastUseReader rep = LastUseReader
  { forall {k} (rep :: k).
LastUseReader rep
-> Op (Aliases rep)
-> Names
-> LastUseM rep (AliasTab, Names, Names)
onOp :: Op (Aliases rep) -> Names -> LastUseM rep (LUTabFun, Names, Names)
  }

type LastUseM rep a = StateT AliasTab (Reader (LastUseReader rep)) a

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

-- | Perform last-use analysis on a 'Prog' in 'SeqMem'
lastUsePrg :: Prog (Aliases SeqMem) -> LUTabPrg
lastUsePrg :: Prog (Aliases SeqMem) -> LUTabPrg
lastUsePrg Prog (Aliases SeqMem)
prg = 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]
map FunDef (Aliases SeqMem) -> (Name, AliasTab)
lastUseSeqMem forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Prog rep -> [FunDef rep]
progFuns Prog (Aliases SeqMem)
prg

-- | Perform last-use analysis on a 'Prog' in 'GPUMem'
lastUsePrgGPU :: Prog (Aliases GPUMem) -> LUTabPrg
lastUsePrgGPU :: Prog (Aliases GPUMem) -> LUTabPrg
lastUsePrgGPU Prog (Aliases GPUMem)
prg = 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]
map FunDef (Aliases GPUMem) -> (Name, AliasTab)
lastUseGPUMem forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Prog rep -> [FunDef rep]
progFuns Prog (Aliases GPUMem)
prg

-- | Perform last-use analysis on a 'FunDef' in 'SeqMem'
lastUseSeqMem :: FunDef (Aliases SeqMem) -> (Name, LUTabFun)
lastUseSeqMem :: FunDef (Aliases SeqMem) -> (Name, AliasTab)
lastUseSeqMem (FunDef Maybe EntryPoint
_ Attrs
_ Name
fname [RetType (Aliases SeqMem)]
_ [FParam (Aliases SeqMem)]
_ Body (Aliases SeqMem)
body) =
  let (AliasTab
res, Names
_) =
        forall r a. Reader r a -> r -> a
runReader
          (forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseBody Body (Aliases SeqMem)
body (forall a. Monoid a => a
mempty, forall a. Monoid a => a
mempty)) forall a. Monoid a => a
mempty)
          (forall {k} (rep :: k).
(Op (Aliases rep)
 -> Names -> LastUseM rep (AliasTab, Names, Names))
-> LastUseReader rep
LastUseReader Op (Aliases SeqMem)
-> Names -> LastUseM SeqMem (AliasTab, Names, Names)
lastUseSeqOp)
   in (Name
fname, AliasTab
res)

-- | Perform last-use analysis on a 'FunDef' in 'GPUMem'
lastUseGPUMem :: FunDef (Aliases GPUMem) -> (Name, LUTabFun)
lastUseGPUMem :: FunDef (Aliases GPUMem) -> (Name, AliasTab)
lastUseGPUMem (FunDef Maybe EntryPoint
_ Attrs
_ Name
fname [RetType (Aliases GPUMem)]
_ [FParam (Aliases GPUMem)]
_ Body (Aliases GPUMem)
body) =
  let (AliasTab
res, Names
_) =
        forall r a. Reader r a -> r -> a
runReader
          (forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseBody Body (Aliases GPUMem)
body (forall a. Monoid a => a
mempty, forall a. Monoid a => a
mempty)) forall a. Monoid a => a
mempty)
          (forall {k} (rep :: k).
(Op (Aliases rep)
 -> Names -> LastUseM rep (AliasTab, Names, Names))
-> LastUseReader rep
LastUseReader Op (Aliases GPUMem)
-> Names -> LastUseM GPUMem (AliasTab, Names, Names)
lastUseGPUOp)
   in (Name
fname, AliasTab
res)

-- | 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 ::
  (ASTRep rep, FreeIn (OpWithAliases (Op 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).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseBody bdy :: Body (Aliases rep)
bdy@(Body BodyDec (Aliases rep)
_ Stms (Aliases rep)
stms Result
result) (AliasTab
lutab, Names
used_nms) = do
  -- perform analysis bottom-up in bindings: results are known to be used,
  -- hence they are added to the used_nms set.
  (AliasTab
lutab', Names
_) <-
    forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Stms (Aliases rep)
-> (AliasTab, Names) -> [VName] -> LastUseM rep (AliasTab, Names)
lastUseStms Stms (Aliases rep)
stms (AliasTab
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 (AliasTab
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 ::
  (CanBeAliased (Op rep), ASTRep 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).
(CanBeAliased (Op rep), ASTRep rep) =>
KernelBody (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseKernelBody bdy :: KernelBody (Aliases rep)
bdy@(KernelBody BodyDec (Aliases rep)
_ Stms (Aliases rep)
stms [KernelResult]
result) (AliasTab
lutab, Names
used_nms) = do
  -- perform analysis bottom-up in bindings: results are known to be used,
  -- hence they are added to the used_nms set.
  (AliasTab
lutab', Names
_) <-
    forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Stms (Aliases rep)
-> (AliasTab, Names) -> [VName] -> LastUseM rep (AliasTab, Names)
lastUseStms Stms (Aliases rep)
stms (AliasTab
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 (AliasTab
lutab', Names
used_nms forall a. Semigroup a => a -> a -> a
<> Names
used_in_body)

lastUseStms ::
  (ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
  Stms (Aliases rep) ->
  (LUTabFun, Names) ->
  [VName] ->
  LastUseM rep (LUTabFun, Names)
lastUseStms :: forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Stms (Aliases rep)
-> (AliasTab, Names) -> [VName] -> LastUseM rep (AliasTab, Names)
lastUseStms Seq (Stm (Aliases rep))
Empty (AliasTab
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 (AliasTab
lutab, Names
nms forall a. Semigroup a => a -> a -> a
<> Names
aliases)
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) (AliasTab
lutab, Names
nms) [VName]
res_nms = 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
  (AliasTab
lutab', Names
nms') <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Stms (Aliases rep)
-> (AliasTab, Names) -> [VName] -> LastUseM rep (AliasTab, Names)
lastUseStms Seq (Stm (Aliases rep))
stms (AliasTab
lutab, Names
nms) [VName]
res_nms
  (AliasTab
lutab'', Names
nms'') <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Stm (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseStm Stm (Aliases rep)
stm (AliasTab
lutab', Names
nms')
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
lutab'', Names
nms'')

lastUseStm ::
  (ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
  Stm (Aliases rep) ->
  (LUTabFun, Names) ->
  LastUseM rep (LUTabFun, Names)
lastUseStm :: forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Stm (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseStm (Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
_ Exp (Aliases rep)
e) (AliasTab
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
    (AliasTab
lutab', Names
last_uses, Names
used_nms') <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Exp (Aliases rep) -> Names -> LastUseM rep (AliasTab, Names, Names)
lastUseExp Exp (Aliases rep)
e 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'' :: AliasTab
lutab'' =
          forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union AliasTab
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 AliasTab
lutab
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
lutab'', Names
used_nms'')

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

-- | Last-Use Analysis for an expression.
lastUseExp ::
  (ASTRep rep, FreeIn (OpWithAliases (Op 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).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Exp (Aliases rep) -> Names -> LastUseM rep (AliasTab, 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.
  (AliasTab
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).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, 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
  (AliasTab
lutab', Names
body_used_nms) <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, 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 (AliasTab
lutab_cases forall a. Semigroup a => a -> a -> a
<> AliasTab
lutab', Names
last_used_arrs, Names
used_nms')
lastUseExp (DoLoop [(FParam (Aliases rep), SubExp)]
var_ses LoopForm (Aliases rep)
_ Body (Aliases rep)
body) Names
used_nms0 = 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)
-> StateT
     AliasTab (Reader (LastUseReader 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)
  (AliasTab
body_lutab, Names
_) <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, 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 :: AliasTab
lutab_res = AliasTab
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 (AliasTab
lutab_res, Names
lu_arrs, Names
used_nms_res)
  where
    initHelper :: Names
-> (Param dec, SubExp)
-> StateT
     AliasTab (Reader (LastUseReader 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 (AliasTab, Names, Names)
on_op <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
reader forall {k} (rep :: k).
LastUseReader rep
-> Op (Aliases rep)
-> Names
-> LastUseM rep (AliasTab, Names, Names)
onOp
  OpWithAliases (Op rep)
-> Names -> LastUseM rep (AliasTab, 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')

lastUseGPUOp :: Op (Aliases GPUMem) -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)
lastUseGPUOp :: Op (Aliases GPUMem)
-> Names -> LastUseM GPUMem (AliasTab, Names, Names)
lastUseGPUOp (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')
lastUseGPUOp (Inner (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 (Inner (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 (Inner (SegOp (SegMap SegLevel
_ SegSpace
_ [Type]
tps KernelBody (Aliases GPUMem)
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
  (AliasTab
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
(CanBeAliased (Op rep), ASTRep rep) =>
KernelBody (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseKernelBody KernelBody (Aliases GPUMem)
kbody (forall a. Monoid a => a
mempty, Names
used_nms')
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
body_lutab, Names
lu_vars, Names
used_nms' forall a. Semigroup a => a -> a -> a
<> Names
used_nms'')
lastUseGPUOp (Inner (SegOp (SegRed SegLevel
_ SegSpace
_ [SegBinOp (Aliases GPUMem)]
sbos [Type]
tps KernelBody (Aliases GPUMem)
kbody))) Names
used_nms = do
  (AliasTab
lutab_sbo, Names
lu_vars_sbo, Names
used_nms_sbo) <- [SegBinOp (Aliases GPUMem)]
-> Names -> LastUseM GPUMem (AliasTab, Names, Names)
lastUseSegBinOp [SegBinOp (Aliases GPUMem)]
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
  (AliasTab
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
(CanBeAliased (Op rep), ASTRep rep) =>
KernelBody (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseKernelBody KernelBody (Aliases GPUMem)
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 AliasTab
lutab_sbo AliasTab
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 (Inner (SegOp (SegScan SegLevel
_ SegSpace
_ [SegBinOp (Aliases GPUMem)]
sbos [Type]
tps KernelBody (Aliases GPUMem)
kbody))) Names
used_nms = do
  (AliasTab
lutab_sbo, Names
lu_vars_sbo, Names
used_nms_sbo) <- [SegBinOp (Aliases GPUMem)]
-> Names -> LastUseM GPUMem (AliasTab, Names, Names)
lastUseSegBinOp [SegBinOp (Aliases GPUMem)]
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
  (AliasTab
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
(CanBeAliased (Op rep), ASTRep rep) =>
KernelBody (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseKernelBody KernelBody (Aliases GPUMem)
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 AliasTab
lutab_sbo AliasTab
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 (Inner (SegOp (SegHist SegLevel
_ SegSpace
_ [HistOp (Aliases GPUMem)]
hos [Type]
tps KernelBody (Aliases GPUMem)
kbody))) Names
used_nms = do
  (AliasTab
lutab_sbo, Names
lu_vars_sbo, Names
used_nms_sbo) <- [HistOp (Aliases GPUMem)]
-> Names -> LastUseM GPUMem (AliasTab, Names, Names)
lastUseHistOp [HistOp (Aliases GPUMem)]
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
  (AliasTab
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
(CanBeAliased (Op rep), ASTRep rep) =>
KernelBody (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseKernelBody KernelBody (Aliases GPUMem)
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 AliasTab
lutab_sbo AliasTab
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 (Inner (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
  (AliasTab
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseBody Body (Aliases GPUMem)
body (forall a. Monoid a => a
mempty, Names
used_nms')
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
body_lutab, Names
lu_vars, Names
used_nms' forall a. Semigroup a => a -> a -> a
<> Names
used_nms'')

lastUseSegBinOp :: [SegBinOp (Aliases GPUMem)] -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)
lastUseSegBinOp :: [SegBinOp (Aliases GPUMem)]
-> Names -> LastUseM GPUMem (AliasTab, Names, Names)
lastUseSegBinOp [SegBinOp (Aliases GPUMem)]
sbos Names
used_nms = do
  ([AliasTab]
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 GPUMem)
-> LastUseM GPUMem (AliasTab, Names, Names)
helper [SegBinOp (Aliases GPUMem)]
sbos
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Monoid a => [a] -> a
mconcat [AliasTab]
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 GPUMem)
-> LastUseM GPUMem (AliasTab, Names, Names)
helper (SegBinOp Commutativity
_ (Lambda [LParam (Aliases GPUMem)]
_ Body (Aliases GPUMem)
body [Type]
_) [SubExp]
neutral Shape
shp) = 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 Shape
shp
      (AliasTab
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseBody Body (Aliases GPUMem)
body (forall a. Monoid a => a
mempty, Names
used_nms')
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
body_lutab, Names
lu_vars, Names
used_nms'')

lastUseHistOp :: [HistOp (Aliases GPUMem)] -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)
lastUseHistOp :: [HistOp (Aliases GPUMem)]
-> Names -> LastUseM GPUMem (AliasTab, Names, Names)
lastUseHistOp [HistOp (Aliases GPUMem)]
hos Names
used_nms = do
  ([AliasTab]
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 GPUMem) -> LastUseM GPUMem (AliasTab, Names, Names)
helper [HistOp (Aliases GPUMem)]
hos
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Monoid a => [a] -> a
mconcat [AliasTab]
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 GPUMem) -> LastUseM GPUMem (AliasTab, Names, Names)
helper (HistOp Shape
shp SubExp
rf [VName]
dest [SubExp]
neutral Shape
shp' (Lambda [LParam (Aliases GPUMem)]
_ Body (Aliases GPUMem)
body [Type]
_)) = 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 Shape
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 Shape
shp'
      (AliasTab
body_lutab, Names
used_nms'') <- forall {k} (rep :: k).
(ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
Body (Aliases rep)
-> (AliasTab, Names) -> LastUseM rep (AliasTab, Names)
lastUseBody Body (Aliases GPUMem)
body (forall a. Monoid a => a
mempty, Names
used_nms')
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (AliasTab
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 (AliasTab, 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 -> StateT AliasTab (Reader (LastUseReader 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 -> StateT AliasTab (Reader (LastUseReader 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'