{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

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

import Control.Monad.Reader
import Data.Bifunctor (first)
import Data.Foldable
import Data.Function ((&))
import Data.Map (Map)
import qualified Data.Map 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 {Env rep -> LastUseOp rep
envLastUseOp :: LastUseOp rep}

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

analyseGPUMem :: Prog GPUMem -> (LastUseMap, Used)
analyseGPUMem :: Prog GPUMem -> (LastUseMap, Used)
analyseGPUMem = LastUseOp GPUMem -> Prog GPUMem -> (LastUseMap, Used)
forall rep inner.
(CanBeAliased (Op rep), Mem rep inner) =>
LastUseOp rep -> Prog rep -> (LastUseMap, Used)
analyseProg LastUseOp GPUMem
analyseGPUOp

analyseSeqMem :: Prog SeqMem -> (LastUseMap, Used)
analyseSeqMem :: Prog SeqMem -> (LastUseMap, Used)
analyseSeqMem = LastUseOp SeqMem -> Prog SeqMem -> (LastUseMap, Used)
forall rep inner.
(CanBeAliased (Op rep), Mem rep inner) =>
LastUseOp rep -> Prog rep -> (LastUseMap, Used)
analyseProg LastUseOp SeqMem
analyseSeqOp

analyseGPUOp :: LastUseOp GPUMem
analyseGPUOp :: LastUseOp GPUMem
analyseGPUOp VName
pat_name (LastUse
lumap, Used
used) (Alloc se sp) = do
  let nms :: Used
nms = (SubExp -> Used
forall a. FreeIn a => a -> Used
freeIn SubExp
se Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Space -> Used
forall a. FreeIn a => a -> Used
freeIn Space
sp) Used -> Used -> Used
`namesSubtract` Used
used
  (LastUse, Used) -> LastUseM GPUMem
forall (m :: * -> *) a. Monad m => a -> m a
return (VName -> Used -> LastUse -> LastUse
insertNames VName
pat_name Used
nms LastUse
lumap, Used
used Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Used
nms)
analyseGPUOp VName
pat_name (LastUse
lumap, Used
used) (Inner (SizeOp sop)) = do
  let nms :: Used
nms = SizeOp -> Used
forall a. FreeIn a => a -> Used
freeIn SizeOp
sop Used -> Used -> Used
`namesSubtract` Used
used
  (LastUse, Used) -> LastUseM GPUMem
forall (m :: * -> *) a. Monad m => a -> m a
return (VName -> Used -> LastUse -> LastUse
insertNames VName
pat_name Used
nms LastUse
lumap, Used
used Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Used
nms)
analyseGPUOp VName
_ (LastUse
lumap, Used
used) (Inner (OtherOp ())) =
  (LastUse, Used) -> LastUseM GPUMem
forall (m :: * -> *) a. Monad m => a -> m a
return (LastUse
lumap, Used
used)
analyseGPUOp VName
pat_name (LastUse
lumap, Used
used) (Inner (SegOp (SegMap lvl _ tps body))) = do
  (LastUse
lumap', Used
used') <- (LastUse, Used) -> KernelBody (Aliases GPUMem) -> LastUseM GPUMem
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Used) -> KernelBody (Aliases rep) -> LastUseM rep
analyseKernelBody (LastUse
lumap, Used
used) KernelBody (Aliases GPUMem)
body
  let nms :: Used
nms = (SegLevel -> Used
forall a. FreeIn a => a -> Used
freeIn SegLevel
lvl Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> [Type] -> Used
forall a. FreeIn a => a -> Used
freeIn [Type]
tps) Used -> Used -> Used
`namesSubtract` Used
used'
  (LastUse, Used) -> LastUseM GPUMem
forall (m :: * -> *) a. Monad m => a -> m a
return (VName -> Used -> LastUse -> LastUse
insertNames VName
pat_name Used
nms LastUse
lumap', Used
used' Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Used
nms)
analyseGPUOp VName
pat_name (LastUse
lumap, Used
used) (Inner (SegOp (SegRed lvl _ binops tps body))) =
  VName
-> LastUse
-> Used
-> SegLevel
-> [SegBinOp (Aliases GPUMem)]
-> [Type]
-> KernelBody (Aliases GPUMem)
-> LastUseM GPUMem
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
VName
-> LastUse
-> Used
-> SegLevel
-> [SegBinOp (Aliases rep)]
-> [Type]
-> KernelBody (Aliases rep)
-> LastUseM rep
segOpHelper VName
pat_name LastUse
lumap Used
used SegLevel
lvl [SegBinOp (Aliases GPUMem)]
binops [Type]
tps KernelBody (Aliases GPUMem)
body
analyseGPUOp VName
pat_name (LastUse
lumap, Used
used) (Inner (SegOp (SegScan lvl _ binops tps body))) =
  VName
-> LastUse
-> Used
-> SegLevel
-> [SegBinOp (Aliases GPUMem)]
-> [Type]
-> KernelBody (Aliases GPUMem)
-> LastUseM GPUMem
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
VName
-> LastUse
-> Used
-> SegLevel
-> [SegBinOp (Aliases rep)]
-> [Type]
-> KernelBody (Aliases rep)
-> LastUseM rep
segOpHelper VName
pat_name LastUse
lumap Used
used SegLevel
lvl [SegBinOp (Aliases GPUMem)]
binops [Type]
tps KernelBody (Aliases GPUMem)
body
analyseGPUOp VName
pat_name (LastUse
lumap, Used
used) (Inner (SegOp (SegHist lvl _ binops tps body))) = do
  (LastUse
lumap', Used
used') <- ((LastUse, Used) -> HistOp (Aliases GPUMem) -> LastUseM GPUMem)
-> (LastUse, Used) -> [HistOp (Aliases GPUMem)] -> LastUseM GPUMem
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (LastUse, Used) -> HistOp (Aliases GPUMem) -> LastUseM GPUMem
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Used) -> HistOp (Aliases rep) -> LastUseM rep
analyseHistOp (LastUse
lumap, Used
used) ([HistOp (Aliases GPUMem)] -> LastUseM GPUMem)
-> [HistOp (Aliases GPUMem)] -> LastUseM GPUMem
forall a b. (a -> b) -> a -> b
$ [HistOp (Aliases GPUMem)] -> [HistOp (Aliases GPUMem)]
forall a. [a] -> [a]
reverse [HistOp (Aliases GPUMem)]
binops
  (LastUse
lumap'', Used
used'') <- (LastUse, Used) -> KernelBody (Aliases GPUMem) -> LastUseM GPUMem
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Used) -> KernelBody (Aliases rep) -> LastUseM rep
analyseKernelBody (LastUse
lumap', Used
used') KernelBody (Aliases GPUMem)
body
  let nms :: Used
nms = (SegLevel -> Used
forall a. FreeIn a => a -> Used
freeIn SegLevel
lvl Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> [Type] -> Used
forall a. FreeIn a => a -> Used
freeIn [Type]
tps) Used -> Used -> Used
`namesSubtract` Used
used''
  (LastUse, Used) -> LastUseM GPUMem
forall (m :: * -> *) a. Monad m => a -> m a
return (VName -> Used -> LastUse -> LastUse
insertNames VName
pat_name Used
nms LastUse
lumap'', Used
used'' Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Used
nms)

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

analyseSeqOp :: LastUseOp SeqMem
analyseSeqOp :: LastUseOp SeqMem
analyseSeqOp VName
pat_name (LastUse
lumap, Used
used) (Alloc se sp) = do
  let nms :: Used
nms = (SubExp -> Used
forall a. FreeIn a => a -> Used
freeIn SubExp
se Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Space -> Used
forall a. FreeIn a => a -> Used
freeIn Space
sp) Used -> Used -> Used
`namesSubtract` Used
used
  (LastUse, Used) -> LastUseM SeqMem
forall (m :: * -> *) a. Monad m => a -> m a
return (VName -> Used -> LastUse -> LastUse
insertNames VName
pat_name Used
nms LastUse
lumap, Used
used Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Used
nms)
analyseSeqOp VName
_ (LastUse
lumap, Used
used) (Inner ()) =
  (LastUse, Used) -> LastUseM SeqMem
forall (m :: * -> *) a. Monad m => a -> m a
return (LastUse
lumap, Used
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 :: LastUseOp rep -> Prog rep -> (LastUseMap, Used)
analyseProg LastUseOp rep
onOp Prog rep
prog =
  Reader (Env rep) (LastUseMap, Used)
-> Env rep -> (LastUseMap, Used)
forall r a. Reader r a -> r -> a
runReader Reader (Env rep) (LastUseMap, Used)
helper (LastUseOp rep -> Env rep
forall rep. LastUseOp rep -> Env rep
Env LastUseOp rep
onOp)
  where
    helper :: Reader (Env rep) (LastUseMap, Used)
helper = do
      let consts :: Used
consts =
            Prog rep -> Stms rep
forall rep. Prog rep -> Stms rep
progConsts Prog rep
prog
              Stms rep -> (Stms rep -> [VName]) -> [VName]
forall a b. a -> (a -> b) -> b
& (Stm rep -> [VName]) -> Stms rep -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([VName] -> [VName]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList ([VName] -> [VName]) -> (Stm rep -> [VName]) -> Stm rep -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (PatElemT (LetDec rep) -> VName)
-> [PatElemT (LetDec rep)] -> [VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap PatElemT (LetDec rep) -> VName
forall dec. PatElemT dec -> VName
patElemName ([PatElemT (LetDec rep)] -> [VName])
-> (Stm rep -> [PatElemT (LetDec rep)]) -> Stm rep -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatT (LetDec rep) -> [PatElemT (LetDec rep)]
forall dec. PatT dec -> [PatElemT dec]
patElems (PatT (LetDec rep) -> [PatElemT (LetDec rep)])
-> (Stm rep -> PatT (LetDec rep))
-> Stm rep
-> [PatElemT (LetDec rep)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm rep -> PatT (LetDec rep)
forall rep. Stm rep -> Pat rep
stmPat)
              [VName] -> ([VName] -> Used) -> Used
forall a b. a -> (a -> b) -> b
& [VName] -> Used
namesFromList
          funs :: [FunDef (Aliases rep)]
funs = Prog (Aliases rep) -> [FunDef (Aliases rep)]
forall rep. Prog rep -> [FunDef rep]
progFuns (Prog (Aliases rep) -> [FunDef (Aliases rep)])
-> Prog (Aliases rep) -> [FunDef (Aliases rep)]
forall a b. (a -> b) -> a -> b
$ Prog rep -> Prog (Aliases rep)
forall rep.
(ASTRep rep, CanBeAliased (Op rep)) =>
Prog rep -> Prog (Aliases rep)
aliasAnalysis Prog rep
prog
      (LastUse
lus, Used
used) <- [(LastUse, Used)] -> (LastUse, Used)
forall a. Monoid a => [a] -> a
mconcat ([(LastUse, Used)] -> (LastUse, Used))
-> ReaderT (Env rep) Identity [(LastUse, Used)]
-> ReaderT (Env rep) Identity (LastUse, Used)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FunDef (Aliases rep)
 -> ReaderT (Env rep) Identity (LastUse, Used))
-> [FunDef (Aliases rep)]
-> ReaderT (Env rep) Identity [(LastUse, Used)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (LastUse
-> Used
-> FunDef (Aliases rep)
-> ReaderT (Env rep) Identity (LastUse, Used)
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Used -> FunDef (Aliases rep) -> LastUseM rep
analyseFun LastUse
forall a. Monoid a => a
mempty Used
consts) [FunDef (Aliases rep)]
funs
      (LastUseMap, Used) -> Reader (Env rep) (LastUseMap, Used)
forall (m :: * -> *) a. Monad m => a -> m a
return (LastUse -> LastUseMap
flipMap LastUse
lus, Used
used)

analyseFun :: (FreeIn (OpWithAliases (Op rep)), ASTRep rep) => LastUse -> Used -> FunDef (Aliases rep) -> LastUseM rep
analyseFun :: LastUse -> Used -> FunDef (Aliases rep) -> LastUseM rep
analyseFun LastUse
lumap Used
used FunDef (Aliases rep)
fun = do
  (LastUse
lumap', Used
used') <- LastUse -> Used -> Body (Aliases rep) -> LastUseM rep
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Used -> Body (Aliases rep) -> LastUseM rep
analyseBody LastUse
lumap Used
used (Body (Aliases rep) -> LastUseM rep)
-> Body (Aliases rep) -> LastUseM rep
forall a b. (a -> b) -> a -> b
$ FunDef (Aliases rep) -> Body (Aliases rep)
forall rep. FunDef rep -> BodyT rep
funDefBody FunDef (Aliases rep)
fun
  (LastUse, Used) -> LastUseM rep
forall (m :: * -> *) a. Monad m => a -> m a
return (LastUse
lumap', Used
used' Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> [Param (FParamInfo rep)] -> Used
forall a. FreeIn a => a -> Used
freeIn (FunDef (Aliases rep) -> [FParam (Aliases rep)]
forall rep. FunDef rep -> [FParam rep]
funDefParams FunDef (Aliases rep)
fun))

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

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

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

    analyseExp :: (LastUse, Used) -> Exp (Aliases rep) -> LastUseM rep
analyseExp (LastUse
lumap, Used
used) (BasicOp BasicOp
_) = do
      let nms :: Used
nms = Exp (Aliases rep) -> Used
forall a. FreeIn a => a -> Used
freeIn Exp (Aliases rep)
e Used -> Used -> Used
`namesSubtract` Used
used
      (LastUse, Used) -> LastUseM rep
forall (m :: * -> *) a. Monad m => a -> m a
return (VName -> Used -> LastUse -> LastUse
insertNames VName
pat_name Used
nms LastUse
lumap, Used
used Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Used
nms)
    analyseExp (LastUse
lumap, Used
used) (Apply Name
_ [(SubExp, Diet)]
args [RetType (Aliases rep)]
_ (Safety, SrcLoc, [SrcLoc])
_) = do
      let nms :: Used
nms = [SubExp] -> Used
forall a. FreeIn a => a -> Used
freeIn ([SubExp] -> Used) -> [SubExp] -> Used
forall a b. (a -> b) -> a -> b
$ ((SubExp, Diet) -> SubExp) -> [(SubExp, Diet)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp, Diet) -> SubExp
forall a b. (a, b) -> a
fst [(SubExp, Diet)]
args
      (LastUse, Used) -> LastUseM rep
forall (m :: * -> *) a. Monad m => a -> m a
return (VName -> Used -> LastUse -> LastUse
insertNames VName
pat_name Used
nms LastUse
lumap, Used
used Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Used
nms)
    analyseExp (LastUse
lumap, Used
used) (If SubExp
cse BodyT (Aliases rep)
then_body BodyT (Aliases rep)
else_body IfDec (BranchType (Aliases rep))
dec) = do
      (LastUse
lumap_then, Used
used_then) <- LastUse -> Used -> BodyT (Aliases rep) -> LastUseM rep
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Used -> Body (Aliases rep) -> LastUseM rep
analyseBody LastUse
lumap Used
used BodyT (Aliases rep)
then_body
      (LastUse
lumap_else, Used
used_else) <- LastUse -> Used -> BodyT (Aliases rep) -> LastUseM rep
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Used -> Body (Aliases rep) -> LastUseM rep
analyseBody LastUse
lumap Used
used BodyT (Aliases rep)
else_body
      let used' :: Used
used' = Used
used_then Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Used
used_else
          nms :: Used
nms = (SubExp -> Used
forall a. FreeIn a => a -> Used
freeIn SubExp
cse Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> IfDec (BranchType rep) -> Used
forall a. FreeIn a => a -> Used
freeIn IfDec (BranchType rep)
IfDec (BranchType (Aliases rep))
dec) Used -> Used -> Used
`namesSubtract` Used
used'
      (LastUse, Used) -> LastUseM rep
forall (m :: * -> *) a. Monad m => a -> m a
return (VName -> Used -> LastUse -> LastUse
insertNames VName
pat_name Used
nms (LastUse
lumap_then LastUse -> LastUse -> LastUse
forall a. Semigroup a => a -> a -> a
<> LastUse
lumap_else), Used
used' Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Used
nms)
    analyseExp (LastUse
lumap, Used
used) (DoLoop [(FParam (Aliases rep), SubExp)]
merge LoopForm (Aliases rep)
form BodyT (Aliases rep)
body) = do
      (LastUse
lumap', Used
used') <- LastUse -> Used -> BodyT (Aliases rep) -> LastUseM rep
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Used -> Body (Aliases rep) -> LastUseM rep
analyseBody LastUse
lumap Used
used BodyT (Aliases rep)
body
      let nms :: Used
nms = ([(Param (FParamInfo rep), SubExp)] -> Used
forall a. FreeIn a => a -> Used
freeIn [(Param (FParamInfo rep), SubExp)]
[(FParam (Aliases rep), SubExp)]
merge Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> LoopForm (Aliases rep) -> Used
forall a. FreeIn a => a -> Used
freeIn LoopForm (Aliases rep)
form) Used -> Used -> Used
`namesSubtract` Used
used'
      (LastUse, Used) -> LastUseM rep
forall (m :: * -> *) a. Monad m => a -> m a
return (VName -> Used -> LastUse -> LastUse
insertNames VName
pat_name Used
nms LastUse
lumap', Used
used' Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Used
nms)
    analyseExp (LastUse
lumap, Used
used) (Op Op (Aliases rep)
op) = do
      VName -> (LastUse, Used) -> OpWithAliases (Op rep) -> LastUseM rep
onOp <- (Env rep
 -> VName
 -> (LastUse, Used)
 -> OpWithAliases (Op rep)
 -> LastUseM rep)
-> ReaderT
     (Env rep)
     Identity
     (VName
      -> (LastUse, Used) -> OpWithAliases (Op rep) -> LastUseM rep)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env rep
-> VName
-> (LastUse, Used)
-> OpWithAliases (Op rep)
-> LastUseM rep
forall rep. Env rep -> LastUseOp rep
envLastUseOp
      VName -> (LastUse, Used) -> OpWithAliases (Op rep) -> LastUseM rep
onOp VName
pat_name (LastUse
lumap, Used
used) Op (Aliases rep)
OpWithAliases (Op rep)
op
    analyseExp (LastUse
lumap, Used
used) (WithAcc [WithAccInput (Aliases rep)]
_ Lambda (Aliases rep)
l) =
      (LastUse, Used) -> Lambda (Aliases rep) -> LastUseM rep
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Used) -> Lambda (Aliases rep) -> LastUseM rep
analyseLambda (LastUse
lumap, Used
used) Lambda (Aliases rep)
l

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

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

analyseSegBinOp ::
  (FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
  (LastUse, Used) ->
  SegBinOp (Aliases rep) ->
  LastUseM rep
analyseSegBinOp :: (LastUse, Used) -> SegBinOp (Aliases rep) -> LastUseM rep
analyseSegBinOp (LastUse
lumap, Used
used) (SegBinOp Commutativity
_ Lambda (Aliases rep)
lambda [SubExp]
neutral Shape
shp) = do
  (LastUse
lumap', Used
used') <- (LastUse, Used) -> Lambda (Aliases rep) -> LastUseM rep
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Used) -> Lambda (Aliases rep) -> LastUseM rep
analyseLambda (LastUse
lumap, Used
used) Lambda (Aliases rep)
lambda
  let nms :: Used
nms = ([SubExp] -> Used
forall a. FreeIn a => a -> Used
freeIn [SubExp]
neutral Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Shape -> Used
forall a. FreeIn a => a -> Used
freeIn Shape
shp) Used -> Used -> Used
`namesSubtract` Used
used'
  (LastUse, Used) -> LastUseM rep
forall (m :: * -> *) a. Monad m => a -> m a
return (LastUse
lumap', Used
used' Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Used
nms)

analyseHistOp ::
  (FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
  (LastUse, Used) ->
  HistOp (Aliases rep) ->
  LastUseM rep
analyseHistOp :: (LastUse, Used) -> HistOp (Aliases rep) -> LastUseM rep
analyseHistOp (LastUse
lumap, Used
used) (HistOp Shape
width SubExp
race [VName]
dest [SubExp]
neutral Shape
shp Lambda (Aliases rep)
lambda) = do
  (LastUse
lumap', Used
used') <- (LastUse, Used) -> Lambda (Aliases rep) -> LastUseM rep
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Used) -> Lambda (Aliases rep) -> LastUseM rep
analyseLambda (LastUse
lumap, Used
used) Lambda (Aliases rep)
lambda
  let nms :: Used
nms =
        ( Shape -> Used
forall a. FreeIn a => a -> Used
freeIn Shape
width Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> SubExp -> Used
forall a. FreeIn a => a -> Used
freeIn SubExp
race Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> [VName] -> Used
forall a. FreeIn a => a -> Used
freeIn [VName]
dest Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> Used
forall a. FreeIn a => a -> Used
freeIn [SubExp]
neutral
            Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Shape -> Used
forall a. FreeIn a => a -> Used
freeIn Shape
shp
        )
          Used -> Used -> Used
`namesSubtract` Used
used'
  (LastUse, Used) -> LastUseM rep
forall (m :: * -> *) a. Monad m => a -> m a
return (LastUse
lumap', Used
used' Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Used
nms)

analyseLambda :: (FreeIn (OpWithAliases (Op rep)), ASTRep rep) => (LastUse, Used) -> Lambda (Aliases rep) -> LastUseM rep
analyseLambda :: (LastUse, Used) -> Lambda (Aliases rep) -> LastUseM rep
analyseLambda (LastUse
lumap, Used
used) (Lambda [LParam (Aliases rep)]
params BodyT (Aliases rep)
body [Type]
ret) = do
  (LastUse
lumap', Used
used') <- LastUse -> Used -> BodyT (Aliases rep) -> LastUseM rep
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Used -> Body (Aliases rep) -> LastUseM rep
analyseBody LastUse
lumap Used
used BodyT (Aliases rep)
body
  let used'' :: Used
used'' = Used
used' Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> [Param (LParamInfo rep)] -> Used
forall a. FreeIn a => a -> Used
freeIn [Param (LParamInfo rep)]
[LParam (Aliases rep)]
params Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> [Type] -> Used
forall a. FreeIn a => a -> Used
freeIn [Type]
ret
  (LastUse, Used) -> LastUseM rep
forall (m :: * -> *) a. Monad m => a -> m a
return (LastUse
lumap', Used
used'')

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

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