{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# 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 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 {forall rep. 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 = LastUseOp GPUMem -> Prog GPUMem -> (Map VName Names, Names)
forall rep 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 = LastUseOp SeqMem -> Prog SeqMem -> (Map VName Names, Names)
forall rep 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 = (SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn SubExp
se Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Space -> Names
forall a. FreeIn a => a -> Names
freeIn Space
sp) Names -> Names -> Names
`namesSubtract` Names
used
  (LastUse, Names) -> LastUseM GPUMem
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap, Names
used Names -> Names -> Names
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 = SizeOp -> Names
forall a. FreeIn a => a -> Names
freeIn SizeOp
sop Names -> Names -> Names
`namesSubtract` Names
used
  (LastUse, Names) -> LastUseM GPUMem
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap, Names
used Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
nms)
analyseGPUOp VName
_ (LastUse
lumap, Names
used) (Inner (OtherOp ())) =
  (LastUse, Names) -> LastUseM GPUMem
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') <- (LastUse, Names) -> KernelBody (Aliases GPUMem) -> LastUseM GPUMem
forall rep.
(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 = (SegLevel -> Names
forall a. FreeIn a => a -> Names
freeIn SegLevel
lvl Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [Type] -> Names
forall a. FreeIn a => a -> Names
freeIn [Type]
tps) Names -> Names -> Names
`namesSubtract` Names
used'
  (LastUse, Names) -> LastUseM GPUMem
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap', Names
used' Names -> Names -> Names
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))) =
  VName
-> LastUse
-> Names
-> SegLevel
-> [SegBinOp (Aliases GPUMem)]
-> [Type]
-> KernelBody (Aliases GPUMem)
-> LastUseM GPUMem
forall rep.
(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))) =
  VName
-> LastUse
-> Names
-> SegLevel
-> [SegBinOp (Aliases GPUMem)]
-> [Type]
-> KernelBody (Aliases GPUMem)
-> LastUseM GPUMem
forall rep.
(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') <- ((LastUse, Names) -> HistOp (Aliases GPUMem) -> LastUseM GPUMem)
-> (LastUse, Names) -> [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, Names) -> HistOp (Aliases GPUMem) -> LastUseM GPUMem
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> HistOp (Aliases rep) -> LastUseM rep
analyseHistOp (LastUse
lumap, Names
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'', Names
used'') <- (LastUse, Names) -> KernelBody (Aliases GPUMem) -> LastUseM GPUMem
forall rep.
(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 = (SegLevel -> Names
forall a. FreeIn a => a -> Names
freeIn SegLevel
lvl Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [Type] -> Names
forall a. FreeIn a => a -> Names
freeIn [Type]
tps) Names -> Names -> Names
`namesSubtract` Names
used''
  (LastUse, Names) -> LastUseM GPUMem
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap'', Names
used'' Names -> Names -> Names
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') <- LastUse -> Names -> Body (Aliases GPUMem) -> LastUseM GPUMem
forall rep.
(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 = [Type] -> Names
forall a. FreeIn a => a -> Names
freeIn [Type]
ts
  (LastUse, Names) -> LastUseM GPUMem
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap', Names
used' Names -> Names -> Names
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 rep.
(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') <- ((LastUse, Names) -> SegBinOp (Aliases rep) -> LastUseM rep)
-> (LastUse, Names) -> [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, Names) -> SegBinOp (Aliases rep) -> LastUseM rep
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> SegBinOp (Aliases rep) -> LastUseM rep
analyseSegBinOp (LastUse
lumap, Names
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'', Names
used'') <- (LastUse, Names) -> KernelBody (Aliases rep) -> LastUseM rep
forall rep.
(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 = (SegLevel -> Names
forall a. FreeIn a => a -> Names
freeIn SegLevel
lvl Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [Type] -> Names
forall a. FreeIn a => a -> Names
freeIn [Type]
tps) Names -> Names -> Names
`namesSubtract` Names
used''
  (LastUse, Names) -> LastUseM rep
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap'', Names
used'' Names -> Names -> Names
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 = (SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn SubExp
se Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Space -> Names
forall a. FreeIn a => a -> Names
freeIn Space
sp) Names -> Names -> Names
`namesSubtract` Names
used
  (LastUse, Names) -> LastUseM SeqMem
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap, Names
used Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
nms)
analyseSeqOp VName
_ (LastUse
lumap, Names
used) (Inner ()) =
  (LastUse, Names) -> LastUseM SeqMem
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 rep inner.
(CanBeAliased (Op rep), Mem rep inner) =>
LastUseOp rep -> Prog rep -> (Map VName Names, Names)
analyseProg LastUseOp rep
onOp Prog rep
prog =
  Reader (Env rep) (Map VName Names, Names)
-> Env rep -> (Map VName Names, Names)
forall r a. Reader r a -> r -> a
runReader Reader (Env rep) (Map VName Names, Names)
helper (LastUseOp rep -> Env rep
forall rep. LastUseOp rep -> Env rep
Env LastUseOp rep
onOp)
  where
    helper :: Reader (Env rep) (Map VName Names, Names)
helper = do
      let consts :: Names
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
. (PatElem (LetDec rep) -> VName)
-> [PatElem (LetDec rep)] -> [VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName ([PatElem (LetDec rep)] -> [VName])
-> (Stm rep -> [PatElem (LetDec rep)]) -> Stm rep -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (LetDec rep) -> [PatElem (LetDec rep)]
forall dec. Pat dec -> [PatElem dec]
patElems (Pat (LetDec rep) -> [PatElem (LetDec rep)])
-> (Stm rep -> Pat (LetDec rep))
-> Stm rep
-> [PatElem (LetDec rep)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm rep -> Pat (LetDec rep)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat)
              [VName] -> ([VName] -> Names) -> Names
forall a b. a -> (a -> b) -> b
& [VName] -> Names
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, Names
used) <- [(LastUse, Names)] -> (LastUse, Names)
forall a. Monoid a => [a] -> a
mconcat ([(LastUse, Names)] -> (LastUse, Names))
-> ReaderT (Env rep) Identity [(LastUse, Names)]
-> ReaderT (Env rep) Identity (LastUse, Names)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FunDef (Aliases rep)
 -> ReaderT (Env rep) Identity (LastUse, Names))
-> [FunDef (Aliases rep)]
-> ReaderT (Env rep) Identity [(LastUse, Names)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (LastUse
-> Names
-> FunDef (Aliases rep)
-> ReaderT (Env rep) Identity (LastUse, Names)
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> FunDef (Aliases rep) -> LastUseM rep
analyseFun LastUse
forall a. Monoid a => a
mempty Names
consts) [FunDef (Aliases rep)]
funs
      (Map VName Names, Names)
-> Reader (Env rep) (Map VName Names, Names)
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 rep.
(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') <- LastUse -> Names -> Body (Aliases rep) -> LastUseM rep
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> Body (Aliases rep) -> LastUseM rep
analyseBody LastUse
lumap Names
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 -> Body rep
funDefBody FunDef (Aliases rep)
fun
  (LastUse, Names) -> LastUseM rep
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LastUse
lumap', Names
used' Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [Param (FParamInfo rep)] -> Names
forall a. FreeIn a => a -> Names
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 :: forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> Stms (Aliases rep) -> LastUseM rep
analyseStms LastUse
lumap Names
used Stms (Aliases rep)
stms = ((LastUse, Names)
 -> Stm (Aliases rep)
 -> ReaderT (Env rep) Identity (LastUse, Names))
-> (LastUse, Names)
-> [Stm (Aliases rep)]
-> ReaderT (Env rep) Identity (LastUse, Names)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (LastUse, Names)
-> Stm (Aliases rep) -> ReaderT (Env rep) Identity (LastUse, Names)
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
(LastUse, Names) -> Stm (Aliases rep) -> LastUseM rep
analyseStm (LastUse
lumap, Names
used) ([Stm (Aliases rep)]
 -> ReaderT (Env rep) Identity (LastUse, Names))
-> [Stm (Aliases rep)]
-> ReaderT (Env rep) Identity (LastUse, Names)
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 :: forall rep.
(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') = Pat (ConsumedInExp, LetDec rep)
-> [PatElem (ConsumedInExp, LetDec rep)]
forall dec. Pat dec -> [PatElem dec]
patElems Pat (ConsumedInExp, LetDec rep)
Pat (LetDec (Aliases rep))
pat [PatElem (ConsumedInExp, LetDec rep)]
-> ([PatElem (ConsumedInExp, LetDec rep)] -> (LastUse, Names))
-> (LastUse, Names)
forall a b. a -> (a -> b) -> b
& ((LastUse, Names)
 -> PatElem (ConsumedInExp, LetDec rep) -> (LastUse, Names))
-> (LastUse, Names)
-> [PatElem (ConsumedInExp, LetDec rep)]
-> (LastUse, Names)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (LastUse, Names)
-> PatElem (ConsumedInExp, LetDec rep) -> (LastUse, Names)
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 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 -> Names -> LastUse -> LastUse
insertNames VName
name' (ConsumedInExp -> Names
unAliases ConsumedInExp
aliases) LastUse
lumap_acc
          Maybe VName
Nothing -> LastUse
lumap_acc,
        Names
used_acc Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> ConsumedInExp -> Names
unAliases ConsumedInExp
aliases
      )

    pat_name :: VName
pat_name = PatElem (ConsumedInExp, LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName (PatElem (ConsumedInExp, LetDec rep) -> VName)
-> PatElem (ConsumedInExp, LetDec rep) -> VName
forall a b. (a -> b) -> a -> b
$ [PatElem (ConsumedInExp, LetDec rep)]
-> PatElem (ConsumedInExp, LetDec rep)
forall a. [a] -> a
head ([PatElem (ConsumedInExp, LetDec rep)]
 -> PatElem (ConsumedInExp, LetDec rep))
-> [PatElem (ConsumedInExp, LetDec rep)]
-> PatElem (ConsumedInExp, LetDec rep)
forall a b. (a -> b) -> a -> b
$ Pat (ConsumedInExp, LetDec rep)
-> [PatElem (ConsumedInExp, LetDec rep)]
forall dec. Pat dec -> [PatElem dec]
patElems Pat (ConsumedInExp, LetDec rep)
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 = Exp (Aliases rep) -> Names
forall a. FreeIn a => a -> Names
freeIn Exp (Aliases rep)
e Names -> Names -> Names
`namesSubtract` Names
used
      (LastUse, Names) -> LastUseM rep
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap, Names
used Names -> Names -> Names
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 = [SubExp] -> Names
forall a. FreeIn a => a -> Names
freeIn ([SubExp] -> Names) -> [SubExp] -> Names
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, Names) -> LastUseM rep
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap, Names
used Names -> Names -> Names
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) <-
        ([LastUse] -> LastUse)
-> ([Names] -> Names) -> ([LastUse], [Names]) -> (LastUse, Names)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap [LastUse] -> LastUse
forall a. Monoid a => [a] -> a
mconcat [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat (([LastUse], [Names]) -> (LastUse, Names))
-> ([(LastUse, Names)] -> ([LastUse], [Names]))
-> [(LastUse, Names)]
-> (LastUse, Names)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(LastUse, Names)] -> ([LastUse], [Names])
forall a b. [(a, b)] -> ([a], [b])
unzip
          ([(LastUse, Names)] -> (LastUse, Names))
-> ReaderT (Env rep) Identity [(LastUse, Names)] -> LastUseM rep
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Case (Body (Aliases rep)) -> LastUseM rep)
-> [Case (Body (Aliases rep))]
-> ReaderT (Env rep) Identity [(LastUse, Names)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (LastUse -> Names -> Body (Aliases rep) -> LastUseM rep
forall rep.
(FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
LastUse -> Names -> Body (Aliases rep) -> LastUseM rep
analyseBody LastUse
lumap Names
used (Body (Aliases rep) -> LastUseM rep)
-> (Case (Body (Aliases rep)) -> Body (Aliases rep))
-> Case (Body (Aliases rep))
-> LastUseM rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Case (Body (Aliases rep)) -> Body (Aliases rep)
forall body. Case body -> body
caseBody) [Case (Body (Aliases rep))]
cases
      (LastUse
lumap_defbody, Names
used_defbody) <- LastUse -> Names -> Body (Aliases rep) -> LastUseM rep
forall rep.
(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 Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
used_defbody
          nms :: Names
nms = ([SubExp] -> Names
forall a. FreeIn a => a -> Names
freeIn [SubExp]
ses Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> MatchDec (BranchType rep) -> Names
forall a. FreeIn a => a -> Names
freeIn MatchDec (BranchType rep)
MatchDec (BranchType (Aliases rep))
dec) Names -> Names -> Names
`namesSubtract` Names
used'
      (LastUse, Names) -> LastUseM rep
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms (LastUse
lumap_cases LastUse -> LastUse -> LastUse
forall a. Semigroup a => a -> a -> a
<> LastUse
lumap_defbody),
          Names
used' Names -> Names -> Names
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') <- LastUse -> Names -> Body (Aliases rep) -> LastUseM rep
forall rep.
(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 = ([(Param (FParamInfo rep), SubExp)] -> Names
forall a. FreeIn a => a -> Names
freeIn [(Param (FParamInfo rep), SubExp)]
[(FParam (Aliases rep), SubExp)]
merge Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> LoopForm (Aliases rep) -> Names
forall a. FreeIn a => a -> Names
freeIn LoopForm (Aliases rep)
form) Names -> Names -> Names
`namesSubtract` Names
used'
      (LastUse, Names) -> LastUseM rep
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms LastUse
lumap', Names
used' Names -> Names -> Names
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 <- (Env rep
 -> VName
 -> (LastUse, Names)
 -> OpWithAliases (Op rep)
 -> LastUseM rep)
-> ReaderT
     (Env rep)
     Identity
     (VName
      -> (LastUse, Names) -> OpWithAliases (Op rep) -> LastUseM rep)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env rep
-> VName
-> (LastUse, Names)
-> OpWithAliases (Op rep)
-> LastUseM rep
forall rep. Env rep -> LastUseOp rep
envLastUseOp
      VName -> (LastUse, Names) -> OpWithAliases (Op rep) -> LastUseM rep
onOp VName
pat_name (LastUse
lumap, Names
used) Op (Aliases rep)
OpWithAliases (Op rep)
op
    analyseExp (LastUse
lumap, Names
used) (WithAcc [WithAccInput (Aliases rep)]
_ Lambda (Aliases rep)
l) =
      (LastUse, Names) -> Lambda (Aliases rep) -> LastUseM rep
forall rep.
(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 rep.
(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 Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Result -> Names
forall a. FreeIn a => a -> Names
freeIn Result
result
  LastUse -> Names -> Stms (Aliases rep) -> LastUseM rep
forall rep.
(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 rep.
(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 Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [KernelResult] -> Names
forall a. FreeIn a => a -> Names
freeIn [KernelResult]
result
   in LastUse -> Names -> Stms (Aliases rep) -> LastUseM rep
forall rep.
(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 rep.
(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') <- (LastUse, Names) -> Lambda (Aliases rep) -> LastUseM rep
forall rep.
(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 = ([SubExp] -> Names
forall a. FreeIn a => a -> Names
freeIn [SubExp]
neutral Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> ShapeBase SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn ShapeBase SubExp
shp) Names -> Names -> Names
`namesSubtract` Names
used'
  (LastUse, Names) -> LastUseM rep
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LastUse
lumap', Names
used' Names -> Names -> Names
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 rep.
(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') <- (LastUse, Names) -> Lambda (Aliases rep) -> LastUseM rep
forall rep.
(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 =
        ( ShapeBase SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn ShapeBase SubExp
width
            Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn SubExp
race
            Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [VName] -> Names
forall a. FreeIn a => a -> Names
freeIn [VName]
dest
            Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> Names
forall a. FreeIn a => a -> Names
freeIn [SubExp]
neutral
            Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> ShapeBase SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn ShapeBase SubExp
shp
        )
          Names -> Names -> Names
`namesSubtract` Names
used'
  (LastUse, Names) -> LastUseM rep
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LastUse
lumap', Names
used' Names -> Names -> Names
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 rep.
(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') <- LastUse -> Names -> Body (Aliases rep) -> LastUseM rep
forall rep.
(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' Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [Param (LParamInfo rep)] -> Names
forall a. FreeIn a => a -> Names
freeIn [Param (LParamInfo rep)]
[LParam (Aliases rep)]
params Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [Type] -> Names
forall a. FreeIn a => a -> Names
freeIn [Type]
ret
  (LastUse, Names) -> LastUseM rep
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 =
  LastUse -> [(VName, VName)]
forall k a. Map k a -> [(k, a)]
M.toList LastUse
m
    [(VName, VName)]
-> ([(VName, VName)] -> [(VName, Names)]) -> [(VName, Names)]
forall a b. a -> (a -> b) -> b
& ((VName, VName) -> (VName, Names))
-> [(VName, VName)] -> [(VName, Names)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Names, VName) -> (VName, Names)
forall a b. (a, b) -> (b, a)
swap ((Names, VName) -> (VName, Names))
-> ((VName, VName) -> (Names, VName))
-> (VName, VName)
-> (VName, Names)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Names) -> (VName, VName) -> (Names, VName)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first VName -> Names
oneName)
    [(VName, Names)]
-> ([(VName, Names)] -> Map VName Names) -> Map VName Names
forall a b. a -> (a -> b) -> b
& ((VName, Names) -> Map VName Names -> Map VName Names)
-> Map VName Names -> [(VName, Names)] -> Map VName Names
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((VName -> Names -> Map VName Names -> Map VName Names)
-> (VName, Names) -> Map VName Names -> Map VName Names
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((VName -> Names -> Map VName Names -> Map VName Names)
 -> (VName, Names) -> Map VName Names -> Map VName Names)
-> (VName -> Names -> Map VName Names -> Map VName Names)
-> (VName, Names)
-> Map VName Names
-> Map VName Names
forall a b. (a -> b) -> a -> b
$ (Names -> Names -> Names)
-> VName -> Names -> Map VName Names -> Map VName Names
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
(<>)) Map VName Names
forall a. Monoid a => a
mempty

insertNames :: VName -> Names -> LastUse -> LastUse
insertNames :: VName -> Names -> LastUse -> LastUse
insertNames VName
name Names
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
$ Names -> [VName]
namesToList Names
names