{-# LANGUAGE TypeFamilies #-}

-- | Provides last-use analysis for Futhark programs.
module Futhark.Analysis.LastUse
  ( LastUseMap,
    LastUse,
    Used,
    analyseGPUMem,
    analyseSeqMem,
  )
where

import Control.Monad.Reader
import Data.Bifunctor (bimap, first)
import Data.Foldable
import Data.Function ((&))
import Data.Map (Map)
import Data.Map qualified as M
import Data.Tuple
import Futhark.Analysis.Alias (aliasAnalysis)
import Futhark.IR.Aliases
import Futhark.IR.GPUMem
import Futhark.IR.SeqMem

-- | 'LastUseMap' tells which names were last used in a given statement.
-- Statements are uniquely identified by the 'VName' of the first value
-- parameter in the statement pattern. 'Names' is the set of names last used.
type LastUseMap = Map VName Names

-- | 'LastUse' is a mapping from a 'VName' to the statement identifying it's
-- last use. 'LastUseMap' is the inverse of 'LastUse'.
type LastUse = Map VName VName

-- | 'Used' is the set of 'VName' that were used somewhere in a
-- statement, body or otherwise.
type Used = Names

type LastUseOp rep =
  VName -> (LastUse, Used) -> Op (Aliases rep) -> LastUseM rep

newtype Env rep = Env {forall {k} (rep :: k). Env rep -> LastUseOp rep
envLastUseOp :: LastUseOp rep}

type LastUseM rep = Reader (Env rep) (LastUse, Used)

analyseGPUMem :: Prog GPUMem -> (LastUseMap, Used)
analyseGPUMem :: Prog GPUMem -> (Map VName Names, Names)
analyseGPUMem = forall {k} (rep :: k) inner.
(CanBeAliased (Op rep), Mem rep inner) =>
LastUseOp rep -> Prog rep -> (Map VName Names, Names)
analyseProg LastUseOp GPUMem
analyseGPUOp

analyseSeqMem :: Prog SeqMem -> (LastUseMap, Used)
analyseSeqMem :: Prog SeqMem -> (Map VName Names, Names)
analyseSeqMem = forall {k} (rep :: k) inner.
(CanBeAliased (Op rep), Mem rep inner) =>
LastUseOp rep -> Prog rep -> (Map VName Names, Names)
analyseProg LastUseOp SeqMem
analyseSeqOp

analyseGPUOp :: LastUseOp GPUMem
analyseGPUOp :: LastUseOp GPUMem
analyseGPUOp VName
pat_name (LastUse
lumap, Names
used) (Alloc SubExp
se Space
sp) = do
  let nms :: Names
nms = (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 -> Names -> Names
`namesSubtract` Names
used
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap, Names
used forall a. Semigroup a => a -> a -> a
<> Names
nms)
analyseGPUOp VName
pat_name (LastUse
lumap, Names
used) (Inner (SizeOp SizeOp
sop)) = do
  let nms :: Names
nms = forall a. FreeIn a => a -> Names
freeIn SizeOp
sop Names -> Names -> Names
`namesSubtract` Names
used
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap, Names
used forall a. Semigroup a => a -> a -> a
<> Names
nms)
analyseGPUOp VName
_ (LastUse
lumap, Names
used) (Inner (OtherOp ())) =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (LastUse
lumap, Names
used)
analyseGPUOp VName
pat_name (LastUse
lumap, Names
used) (Inner (SegOp (SegMap SegLevel
lvl SegSpace
_ [Type]
tps KernelBody (Aliases GPUMem)
body))) = do
  (LastUse
lumap', Names
used') <- forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> KernelBody (Aliases rep) -> LastUseM rep
analyseKernelBody (LastUse
lumap, Names
used) KernelBody (Aliases GPUMem)
body
  let nms :: Names
nms = (forall a. FreeIn a => a -> Names
freeIn SegLevel
lvl forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn [Type]
tps) Names -> Names -> Names
`namesSubtract` Names
used'
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap', Names
used' forall a. Semigroup a => a -> a -> a
<> Names
nms)
analyseGPUOp VName
pat_name (LastUse
lumap, Names
used) (Inner (SegOp (SegRed SegLevel
lvl SegSpace
_ [SegBinOp (Aliases GPUMem)]
binops [Type]
tps KernelBody (Aliases GPUMem)
body))) =
  forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
VName
-> LastUse
-> Names
-> SegLevel
-> [SegBinOp (Aliases rep)]
-> [Type]
-> KernelBody (Aliases rep)
-> LastUseM rep
segOpHelper VName
pat_name LastUse
lumap Names
used SegLevel
lvl [SegBinOp (Aliases GPUMem)]
binops [Type]
tps KernelBody (Aliases GPUMem)
body
analyseGPUOp VName
pat_name (LastUse
lumap, Names
used) (Inner (SegOp (SegScan SegLevel
lvl SegSpace
_ [SegBinOp (Aliases GPUMem)]
binops [Type]
tps KernelBody (Aliases GPUMem)
body))) =
  forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
VName
-> LastUse
-> Names
-> SegLevel
-> [SegBinOp (Aliases rep)]
-> [Type]
-> KernelBody (Aliases rep)
-> LastUseM rep
segOpHelper VName
pat_name LastUse
lumap Names
used SegLevel
lvl [SegBinOp (Aliases GPUMem)]
binops [Type]
tps KernelBody (Aliases GPUMem)
body
analyseGPUOp VName
pat_name (LastUse
lumap, Names
used) (Inner (SegOp (SegHist SegLevel
lvl SegSpace
_ [HistOp (Aliases GPUMem)]
binops [Type]
tps KernelBody (Aliases GPUMem)
body))) = do
  (LastUse
lumap', Names
used') <- forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> HistOp (Aliases rep) -> LastUseM rep
analyseHistOp (LastUse
lumap, Names
used) forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [a]
reverse [HistOp (Aliases GPUMem)]
binops
  (LastUse
lumap'', Names
used'') <- forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> KernelBody (Aliases rep) -> LastUseM rep
analyseKernelBody (LastUse
lumap', Names
used') KernelBody (Aliases GPUMem)
body
  let nms :: Names
nms = (forall a. FreeIn a => a -> Names
freeIn SegLevel
lvl forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn [Type]
tps) Names -> Names -> Names
`namesSubtract` Names
used''
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap'', Names
used'' forall a. Semigroup a => a -> a -> a
<> Names
nms)
analyseGPUOp VName
pat_name (LastUse
lumap, Names
used) (Inner (GPUBody [Type]
ts Body (Aliases GPUMem)
body)) = do
  (LastUse
lumap', Names
used') <- forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> Body (Aliases rep) -> LastUseM rep
analyseBody LastUse
lumap Names
used Body (Aliases GPUMem)
body
  let nms :: Names
nms = forall a. FreeIn a => a -> Names
freeIn [Type]
ts
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap', Names
used' forall a. Semigroup a => a -> a -> a
<> Names
nms)

segOpHelper ::
  (FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
  VName ->
  LastUse ->
  Used ->
  SegLevel ->
  [SegBinOp (Aliases rep)] ->
  [Type] ->
  KernelBody (Aliases rep) ->
  LastUseM rep
segOpHelper :: forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
VName
-> LastUse
-> Names
-> SegLevel
-> [SegBinOp (Aliases rep)]
-> [Type]
-> KernelBody (Aliases rep)
-> LastUseM rep
segOpHelper VName
pat_name LastUse
lumap Names
used SegLevel
lvl [SegBinOp (Aliases rep)]
binops [Type]
tps KernelBody (Aliases rep)
body = do
  (LastUse
lumap', Names
used') <- forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> SegBinOp (Aliases rep) -> LastUseM rep
analyseSegBinOp (LastUse
lumap, Names
used) forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [a]
reverse [SegBinOp (Aliases rep)]
binops
  (LastUse
lumap'', Names
used'') <- forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> KernelBody (Aliases rep) -> LastUseM rep
analyseKernelBody (LastUse
lumap', Names
used') KernelBody (Aliases rep)
body
  let nms :: Names
nms = (forall a. FreeIn a => a -> Names
freeIn SegLevel
lvl forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn [Type]
tps) Names -> Names -> Names
`namesSubtract` Names
used''
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap'', Names
used'' forall a. Semigroup a => a -> a -> a
<> Names
nms)

analyseSeqOp :: LastUseOp SeqMem
analyseSeqOp :: LastUseOp SeqMem
analyseSeqOp VName
pat_name (LastUse
lumap, Names
used) (Alloc SubExp
se Space
sp) = do
  let nms :: Names
nms = (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 -> Names -> Names
`namesSubtract` Names
used
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap, Names
used forall a. Semigroup a => a -> a -> a
<> Names
nms)
analyseSeqOp VName
_ (LastUse
lumap, Names
used) (Inner ()) =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (LastUse
lumap, Names
used)

-- | Analyses a program to return a last-use map, mapping each simple statement
-- in the program to the values that were last used within that statement, and
-- the set of all `VName` that were used inside.
analyseProg :: (CanBeAliased (Op rep), Mem rep inner) => LastUseOp rep -> Prog rep -> (LastUseMap, Used)
analyseProg :: forall {k} (rep :: k) inner.
(CanBeAliased (Op rep), Mem rep inner) =>
LastUseOp rep -> Prog rep -> (Map VName Names, Names)
analyseProg LastUseOp rep
onOp Prog rep
prog =
  forall r a. Reader r a -> r -> a
runReader ReaderT (Env rep) Identity (Map VName Names, Names)
helper (forall {k} (rep :: k). LastUseOp rep -> Env rep
Env LastUseOp rep
onOp)
  where
    helper :: ReaderT (Env rep) Identity (Map VName Names, Names)
helper = do
      let consts :: Names
consts =
            forall {k} (rep :: k). Prog rep -> Stms rep
progConsts Prog rep
prog
              forall a b. a -> (a -> b) -> b
& forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall dec. PatElem dec -> VName
patElemName forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. Pat dec -> [PatElem dec]
patElems 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
          funs :: [FunDef (Aliases rep)]
funs = forall {k} (rep :: k). Prog rep -> [FunDef rep]
progFuns forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
(ASTRep rep, CanBeAliased (Op rep)) =>
Prog rep -> Prog (Aliases rep)
aliasAnalysis Prog rep
prog
      (LastUse
lus, Names
used) <- forall a. Monoid a => [a] -> a
mconcat 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).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> FunDef (Aliases rep) -> LastUseM rep
analyseFun forall a. Monoid a => a
mempty Names
consts) [FunDef (Aliases rep)]
funs
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (LastUse -> Map VName Names
flipMap LastUse
lus, Names
used)

analyseFun :: (FreeIn (OpWithAliases (Op rep)), ASTRep rep) => LastUse -> Used -> FunDef (Aliases rep) -> LastUseM rep
analyseFun :: forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> FunDef (Aliases rep) -> LastUseM rep
analyseFun LastUse
lumap Names
used FunDef (Aliases rep)
fun = do
  (LastUse
lumap', Names
used') <- forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> Body (Aliases rep) -> LastUseM rep
analyseBody LastUse
lumap Names
used forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). FunDef rep -> Body rep
funDefBody FunDef (Aliases rep)
fun
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (LastUse
lumap', Names
used' forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn (forall {k} (rep :: k). FunDef rep -> [FParam rep]
funDefParams FunDef (Aliases rep)
fun))

analyseStms :: (FreeIn (OpWithAliases (Op rep)), ASTRep rep) => LastUse -> Used -> Stms (Aliases rep) -> LastUseM rep
analyseStms :: forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> Stms (Aliases rep) -> LastUseM rep
analyseStms LastUse
lumap Names
used Stms (Aliases rep)
stms = forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> Stm (Aliases rep) -> LastUseM rep
analyseStm (LastUse
lumap, Names
used) forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [a]
reverse forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Stms rep -> [Stm rep]
stmsToList Stms (Aliases rep)
stms

analyseStm :: (FreeIn (OpWithAliases (Op rep)), ASTRep rep) => (LastUse, Used) -> Stm (Aliases rep) -> LastUseM rep
analyseStm :: forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> Stm (Aliases rep) -> LastUseM rep
analyseStm (LastUse
lumap0, Names
used0) (Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
_ Exp (Aliases rep)
e) = do
  let (LastUse
lumap', Names
used') = forall dec. Pat dec -> [PatElem dec]
patElems Pat (LetDec (Aliases rep))
pat forall a b. a -> (a -> b) -> b
& forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall {b}.
(LastUse, Names) -> PatElem (ConsumedInExp, b) -> (LastUse, Names)
helper (LastUse
lumap0, Names
used0)
  (LastUse, Names) -> Exp (Aliases rep) -> LastUseM rep
analyseExp (LastUse
lumap', Names
used') Exp (Aliases rep)
e
  where
    helper :: (LastUse, Names) -> PatElem (ConsumedInExp, b) -> (LastUse, Names)
helper (LastUse
lumap_acc, Names
used_acc) (PatElem VName
name (ConsumedInExp
aliases, b
_)) =
      -- Any aliases of `name` should have the same last-use as `name`
      ( case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name LastUse
lumap_acc of
          Just VName
name' ->
            VName -> Names -> LastUse -> LastUse
insertNames VName
name' (ConsumedInExp -> Names
unAliases ConsumedInExp
aliases) LastUse
lumap_acc
          Maybe VName
Nothing -> LastUse
lumap_acc,
        Names
used_acc forall a. Semigroup a => a -> a -> a
<> ConsumedInExp -> Names
unAliases ConsumedInExp
aliases
      )

    pat_name :: VName
pat_name = forall dec. PatElem dec -> VName
patElemName forall a b. (a -> b) -> a -> b
$ forall a. [a] -> a
head forall a b. (a -> b) -> a -> b
$ forall dec. Pat dec -> [PatElem dec]
patElems Pat (LetDec (Aliases rep))
pat

    analyseExp :: (LastUse, Names) -> Exp (Aliases rep) -> LastUseM rep
analyseExp (LastUse
lumap, Names
used) (BasicOp BasicOp
_) = do
      let nms :: Names
nms = forall a. FreeIn a => a -> Names
freeIn Exp (Aliases rep)
e Names -> Names -> Names
`namesSubtract` Names
used
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap, Names
used forall a. Semigroup a => a -> a -> a
<> Names
nms)
    analyseExp (LastUse
lumap, Names
used) (Apply Name
_ [(SubExp, Diet)]
args [RetType (Aliases rep)]
_ (Safety, SrcLoc, [SrcLoc])
_) = do
      let nms :: Names
nms = forall a. FreeIn a => a -> Names
freeIn forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(SubExp, Diet)]
args
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap, Names
used forall a. Semigroup a => a -> a -> a
<> Names
nms)
    analyseExp (LastUse
lumap, Names
used) (Match [SubExp]
ses [Case (Body (Aliases rep))]
cases Body (Aliases rep)
defbody MatchDec (BranchType (Aliases rep))
dec) = do
      (LastUse
lumap_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 {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> Body (Aliases rep) -> LastUseM rep
analyseBody LastUse
lumap Names
used forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall body. Case body -> body
caseBody) [Case (Body (Aliases rep))]
cases
      (LastUse
lumap_defbody, Names
used_defbody) <- forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> Body (Aliases rep) -> LastUseM rep
analyseBody LastUse
lumap Names
used Body (Aliases rep)
defbody
      let used' :: Names
used' = Names
used_cases forall a. Semigroup a => a -> a -> a
<> Names
used_defbody
          nms :: Names
nms = (forall a. FreeIn a => a -> Names
freeIn [SubExp]
ses forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn MatchDec (BranchType (Aliases rep))
dec) Names -> Names -> Names
`namesSubtract` Names
used'
      forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms (LastUse
lumap_cases forall a. Semigroup a => a -> a -> a
<> LastUse
lumap_defbody),
          Names
used' forall a. Semigroup a => a -> a -> a
<> Names
nms
        )
    analyseExp (LastUse
lumap, Names
used) (DoLoop [(FParam (Aliases rep), SubExp)]
merge LoopForm (Aliases rep)
form Body (Aliases rep)
body) = do
      (LastUse
lumap', Names
used') <- forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> Body (Aliases rep) -> LastUseM rep
analyseBody LastUse
lumap Names
used Body (Aliases rep)
body
      let nms :: Names
nms = (forall a. FreeIn a => a -> Names
freeIn [(FParam (Aliases rep), SubExp)]
merge forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn LoopForm (Aliases rep)
form) Names -> Names -> Names
`namesSubtract` Names
used'
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap', Names
used' forall a. Semigroup a => a -> a -> a
<> Names
nms)
    analyseExp (LastUse
lumap, Names
used) (Op Op (Aliases rep)
op) = do
      VName -> (LastUse, Names) -> OpWithAliases (Op rep) -> LastUseM rep
onOp <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall {k} (rep :: k). Env rep -> LastUseOp rep
envLastUseOp
      VName -> (LastUse, Names) -> OpWithAliases (Op rep) -> LastUseM rep
onOp VName
pat_name (LastUse
lumap, Names
used) Op (Aliases rep)
op
    analyseExp (LastUse
lumap, Names
used) (WithAcc [WithAccInput (Aliases rep)]
_ Lambda (Aliases rep)
l) =
      forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> Lambda (Aliases rep) -> LastUseM rep
analyseLambda (LastUse
lumap, Names
used) Lambda (Aliases rep)
l

analyseBody :: (FreeIn (OpWithAliases (Op rep)), ASTRep rep) => LastUse -> Used -> Body (Aliases rep) -> LastUseM rep
analyseBody :: forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> Body (Aliases rep) -> LastUseM rep
analyseBody LastUse
lumap Names
used (Body BodyDec (Aliases rep)
_ Stms (Aliases rep)
stms Result
result) = do
  let used' :: Names
used' = Names
used forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn Result
result
  forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> Stms (Aliases rep) -> LastUseM rep
analyseStms LastUse
lumap Names
used' Stms (Aliases rep)
stms

analyseKernelBody ::
  (FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
  (LastUse, Used) ->
  KernelBody (Aliases rep) ->
  LastUseM rep
analyseKernelBody :: forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> KernelBody (Aliases rep) -> LastUseM rep
analyseKernelBody (LastUse
lumap, Names
used) (KernelBody BodyDec (Aliases rep)
_ Stms (Aliases rep)
stms [KernelResult]
result) =
  let used' :: Names
used' = Names
used forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn [KernelResult]
result
   in forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> Stms (Aliases rep) -> LastUseM rep
analyseStms LastUse
lumap Names
used' Stms (Aliases rep)
stms

analyseSegBinOp ::
  (FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
  (LastUse, Used) ->
  SegBinOp (Aliases rep) ->
  LastUseM rep
analyseSegBinOp :: forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> SegBinOp (Aliases rep) -> LastUseM rep
analyseSegBinOp (LastUse
lumap, Names
used) (SegBinOp Commutativity
_ Lambda (Aliases rep)
lambda [SubExp]
neutral ShapeBase SubExp
shp) = do
  (LastUse
lumap', Names
used') <- forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> Lambda (Aliases rep) -> LastUseM rep
analyseLambda (LastUse
lumap, Names
used) Lambda (Aliases rep)
lambda
  let nms :: Names
nms = (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) Names -> Names -> Names
`namesSubtract` Names
used'
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (LastUse
lumap', Names
used' forall a. Semigroup a => a -> a -> a
<> Names
nms)

analyseHistOp ::
  (FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
  (LastUse, Used) ->
  HistOp (Aliases rep) ->
  LastUseM rep
analyseHistOp :: forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> HistOp (Aliases rep) -> LastUseM rep
analyseHistOp (LastUse
lumap, Names
used) (HistOp ShapeBase SubExp
width SubExp
race [VName]
dest [SubExp]
neutral ShapeBase SubExp
shp Lambda (Aliases rep)
lambda) = do
  (LastUse
lumap', Names
used') <- forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> Lambda (Aliases rep) -> LastUseM rep
analyseLambda (LastUse
lumap, Names
used) Lambda (Aliases rep)
lambda
  let nms :: Names
nms =
        ( forall a. FreeIn a => a -> Names
freeIn ShapeBase SubExp
width
            forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn SubExp
race
            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
        )
          Names -> Names -> Names
`namesSubtract` Names
used'
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (LastUse
lumap', Names
used' forall a. Semigroup a => a -> a -> a
<> Names
nms)

analyseLambda :: (FreeIn (OpWithAliases (Op rep)), ASTRep rep) => (LastUse, Used) -> Lambda (Aliases rep) -> LastUseM rep
analyseLambda :: forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> Lambda (Aliases rep) -> LastUseM rep
analyseLambda (LastUse
lumap, Names
used) (Lambda [LParam (Aliases rep)]
params Body (Aliases rep)
body [Type]
ret) = do
  (LastUse
lumap', Names
used') <- forall {k} (rep :: k).
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> Body (Aliases rep) -> LastUseM rep
analyseBody LastUse
lumap Names
used Body (Aliases rep)
body
  let used'' :: Names
used'' = Names
used' forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn [LParam (Aliases rep)]
params forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn [Type]
ret
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (LastUse
lumap', Names
used'')

flipMap :: Map VName VName -> Map VName Names
flipMap :: LastUse -> Map VName Names
flipMap LastUse
m =
  forall k a. Map k a -> [(k, a)]
M.toList LastUse
m
    forall a b. a -> (a -> b) -> b
& forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a, b) -> (b, a)
swap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first VName -> Names
oneName)
    forall a b. a -> (a -> b) -> b
& forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith forall a. Semigroup a => a -> a -> a
(<>)) forall a. Monoid a => a
mempty

insertNames :: VName -> Names -> LastUse -> LastUse
insertNames :: VName -> Names -> LastUse -> LastUse
insertNames VName
name Names
names LastUse
lumap =
  forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith forall a b. (a -> b) -> a -> b
$ \VName
_ VName
x -> VName
x) VName
name) LastUse
lumap forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList Names
names