{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}

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

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

-- | `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

-- | 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 :: Prog GPUMem -> (LastUseMap, Used)
analyseProg :: Prog GPUMem -> (LastUseMap, Used)
analyseProg Prog GPUMem
prog =
  let consts :: Used
consts =
        Prog GPUMem -> Stms GPUMem
forall rep. Prog rep -> Stms rep
progConsts Prog GPUMem
prog
          Stms GPUMem -> (Stms GPUMem -> [VName]) -> [VName]
forall a b. a -> (a -> b) -> b
& (Stm GPUMem -> [VName]) -> Stms GPUMem -> [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 GPUMem -> [VName]) -> Stm GPUMem -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (PatElemT LetDecMem -> VName) -> [PatElemT LetDecMem] -> [VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap PatElemT LetDecMem -> VName
forall dec. PatElemT dec -> VName
patElemName ([PatElemT LetDecMem] -> [VName])
-> (Stm GPUMem -> [PatElemT LetDecMem]) -> Stm GPUMem -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatT LetDecMem -> [PatElemT LetDecMem]
forall dec. PatT dec -> [PatElemT dec]
patElems (PatT LetDecMem -> [PatElemT LetDecMem])
-> (Stm GPUMem -> PatT LetDecMem)
-> Stm GPUMem
-> [PatElemT LetDecMem]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm GPUMem -> PatT LetDecMem
forall rep. Stm rep -> Pat rep
stmPat)
          [VName] -> ([VName] -> Used) -> Used
forall a b. a -> (a -> b) -> b
& [VName] -> Used
namesFromList
      funs :: [FunDef (Aliases GPUMem)]
funs = Prog (Aliases GPUMem) -> [FunDef (Aliases GPUMem)]
forall rep. Prog rep -> [FunDef rep]
progFuns (Prog (Aliases GPUMem) -> [FunDef (Aliases GPUMem)])
-> Prog (Aliases GPUMem) -> [FunDef (Aliases GPUMem)]
forall a b. (a -> b) -> a -> b
$ Prog GPUMem -> Prog (Aliases GPUMem)
forall rep.
(ASTRep rep, CanBeAliased (Op rep)) =>
Prog rep -> Prog (Aliases rep)
aliasAnalysis Prog GPUMem
prog
      (LastUse
lus, Used
used) = (FunDef (Aliases GPUMem) -> (LastUse, Used))
-> [FunDef (Aliases GPUMem)] -> (LastUse, Used)
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (LastUse -> Used -> FunDef (Aliases GPUMem) -> (LastUse, Used)
analyseFun LastUse
forall a. Monoid a => a
mempty Used
consts) [FunDef (Aliases GPUMem)]
funs
   in (LastUse -> LastUseMap
flipMap LastUse
lus, Used
used)

analyseFun :: LastUse -> Used -> FunDef (Aliases GPUMem) -> (LastUse, Used)
analyseFun :: LastUse -> Used -> FunDef (Aliases GPUMem) -> (LastUse, Used)
analyseFun LastUse
lumap Used
used FunDef (Aliases GPUMem)
fun =
  let (LastUse
lumap', Used
used') = LastUse -> Used -> Body (Aliases GPUMem) -> (LastUse, Used)
analyseBody LastUse
lumap Used
used (Body (Aliases GPUMem) -> (LastUse, Used))
-> Body (Aliases GPUMem) -> (LastUse, Used)
forall a b. (a -> b) -> a -> b
$ FunDef (Aliases GPUMem) -> Body (Aliases GPUMem)
forall rep. FunDef rep -> BodyT rep
funDefBody FunDef (Aliases GPUMem)
fun
   in (LastUse
lumap', Used
used' Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> [Param FParamMem] -> Used
forall a. FreeIn a => a -> Used
freeIn (FunDef (Aliases GPUMem) -> [FParam (Aliases GPUMem)]
forall rep. FunDef rep -> [FParam rep]
funDefParams FunDef (Aliases GPUMem)
fun))

analyseStms :: LastUse -> Used -> Stms (Aliases GPUMem) -> (LastUse, Used)
analyseStms :: LastUse -> Used -> Stms (Aliases GPUMem) -> (LastUse, Used)
analyseStms LastUse
lumap Used
used Stms (Aliases GPUMem)
stms = (Stm (Aliases GPUMem) -> (LastUse, Used) -> (LastUse, Used))
-> (LastUse, Used) -> [Stm (Aliases GPUMem)] -> (LastUse, Used)
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Stm (Aliases GPUMem) -> (LastUse, Used) -> (LastUse, Used)
analyseStm (LastUse
lumap, Used
used) ([Stm (Aliases GPUMem)] -> (LastUse, Used))
-> [Stm (Aliases GPUMem)] -> (LastUse, Used)
forall a b. (a -> b) -> a -> b
$ Stms (Aliases GPUMem) -> [Stm (Aliases GPUMem)]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms (Aliases GPUMem)
stms

analyseStm :: Stm (Aliases GPUMem) -> (LastUse, Used) -> (LastUse, Used)
analyseStm :: Stm (Aliases GPUMem) -> (LastUse, Used) -> (LastUse, Used)
analyseStm (Let Pat (Aliases GPUMem)
pat StmAux (ExpDec (Aliases GPUMem))
_ Exp (Aliases GPUMem)
e) (LastUse
lumap0, Used
used0) =
  let (LastUse
lumap', Used
used') = PatT (AliasDec, LetDecMem) -> [PatElemT (AliasDec, LetDecMem)]
forall dec. PatT dec -> [PatElemT dec]
patElems PatT (AliasDec, LetDecMem)
Pat (Aliases GPUMem)
pat [PatElemT (AliasDec, LetDecMem)]
-> ([PatElemT (AliasDec, LetDecMem)] -> (LastUse, Used))
-> (LastUse, Used)
forall a b. a -> (a -> b) -> b
& ((LastUse, Used)
 -> PatElemT (AliasDec, LetDecMem) -> (LastUse, Used))
-> (LastUse, Used)
-> [PatElemT (AliasDec, LetDecMem)]
-> (LastUse, Used)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (LastUse, Used)
-> PatElemT (AliasDec, LetDecMem) -> (LastUse, Used)
forall b.
(LastUse, Used) -> PatElemT (AliasDec, b) -> (LastUse, Used)
helper (LastUse
lumap0, Used
used0)
   in (LastUse, Used) -> Exp (Aliases GPUMem) -> (LastUse, Used)
analyseExp (LastUse
lumap', Used
used') Exp (Aliases GPUMem)
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, LetDecMem) -> VName
forall dec. PatElemT dec -> VName
patElemName (PatElemT (AliasDec, LetDecMem) -> VName)
-> PatElemT (AliasDec, LetDecMem) -> VName
forall a b. (a -> b) -> a -> b
$ [PatElemT (AliasDec, LetDecMem)] -> PatElemT (AliasDec, LetDecMem)
forall a. [a] -> a
head ([PatElemT (AliasDec, LetDecMem)]
 -> PatElemT (AliasDec, LetDecMem))
-> [PatElemT (AliasDec, LetDecMem)]
-> PatElemT (AliasDec, LetDecMem)
forall a b. (a -> b) -> a -> b
$ PatT (AliasDec, LetDecMem) -> [PatElemT (AliasDec, LetDecMem)]
forall dec. PatT dec -> [PatElemT dec]
patElems PatT (AliasDec, LetDecMem)
Pat (Aliases GPUMem)
pat
    analyseExp :: (LastUse, Used) -> Exp (Aliases GPUMem) -> (LastUse, Used)
    analyseExp :: (LastUse, Used) -> Exp (Aliases GPUMem) -> (LastUse, Used)
analyseExp (LastUse
lumap, Used
used) (BasicOp BasicOp
_) =
      let nms :: Used
nms = Exp (Aliases GPUMem) -> Used
forall a. FreeIn a => a -> Used
freeIn Exp (Aliases GPUMem)
e Used -> Used -> Used
`namesSubtract` Used
used
       in (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 GPUMem)]
_ (Safety, SrcLoc, [SrcLoc])
_) =
      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
       in (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 Body (Aliases GPUMem)
then_body Body (Aliases GPUMem)
else_body IfDec (BranchType (Aliases GPUMem))
dec) =
      let (LastUse
lumap_then, Used
used_then) = LastUse -> Used -> Body (Aliases GPUMem) -> (LastUse, Used)
analyseBody LastUse
lumap Used
used Body (Aliases GPUMem)
then_body
          (LastUse
lumap_else, Used
used_else) = LastUse -> Used -> Body (Aliases GPUMem) -> (LastUse, Used)
analyseBody LastUse
lumap Used
used Body (Aliases GPUMem)
else_body
          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 BranchTypeMem -> Used
forall a. FreeIn a => a -> Used
freeIn IfDec (BranchType (Aliases GPUMem))
IfDec BranchTypeMem
dec) Used -> Used -> Used
`namesSubtract` Used
used')
       in (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 GPUMem), SubExp)]
merge LoopForm (Aliases GPUMem)
form Body (Aliases GPUMem)
body) =
      let (LastUse
lumap', Used
used') = LastUse -> Used -> Body (Aliases GPUMem) -> (LastUse, Used)
analyseBody LastUse
lumap Used
used Body (Aliases GPUMem)
body
          nms :: Used
nms = ([(Param FParamMem, SubExp)] -> Used
forall a. FreeIn a => a -> Used
freeIn [(FParam (Aliases GPUMem), SubExp)]
[(Param FParamMem, SubExp)]
merge Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> LoopForm (Aliases GPUMem) -> Used
forall a. FreeIn a => a -> Used
freeIn LoopForm (Aliases GPUMem)
form) Used -> Used -> Used
`namesSubtract` Used
used'
       in (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 (Alloc se sp)) =
      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
       in (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 (Inner (SizeOp sop))) =
      let nms :: Used
nms = SizeOp -> Used
forall a. FreeIn a => a -> Used
freeIn SizeOp
sop Used -> Used -> Used
`namesSubtract` Used
used
       in (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 (Inner (OtherOp ()))) =
      (LastUse
lumap, Used
used)
    analyseExp (LastUse
lumap, Used
used) (Op (Inner (SegOp (SegMap lvl _ tps body)))) =
      let (LastUse
lumap', Used
used') = (LastUse, Used) -> KernelBody (Aliases GPUMem) -> (LastUse, Used)
analyseKernelBody (LastUse
lumap, Used
used) KernelBody (Aliases GPUMem)
body
          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'
       in (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 (Inner (SegOp (SegRed lvl _ binops tps body)))) =
      LastUse
-> Used
-> SegLevel
-> [SegBinOp (Aliases GPUMem)]
-> [Type]
-> KernelBody (Aliases GPUMem)
-> (LastUse, Used)
forall (t :: * -> *) a a.
(Foldable t, FreeIn a, FreeIn a) =>
LastUse
-> Used
-> a
-> t (SegBinOp (Aliases GPUMem))
-> a
-> KernelBody (Aliases GPUMem)
-> (LastUse, Used)
segOpHelper LastUse
lumap Used
used SegLevel
lvl [SegBinOp (Aliases GPUMem)]
binops [Type]
tps KernelBody (Aliases GPUMem)
body
    analyseExp (LastUse
lumap, Used
used) (Op (Inner (SegOp (SegScan lvl _ binops tps body)))) =
      LastUse
-> Used
-> SegLevel
-> [SegBinOp (Aliases GPUMem)]
-> [Type]
-> KernelBody (Aliases GPUMem)
-> (LastUse, Used)
forall (t :: * -> *) a a.
(Foldable t, FreeIn a, FreeIn a) =>
LastUse
-> Used
-> a
-> t (SegBinOp (Aliases GPUMem))
-> a
-> KernelBody (Aliases GPUMem)
-> (LastUse, Used)
segOpHelper LastUse
lumap Used
used SegLevel
lvl [SegBinOp (Aliases GPUMem)]
binops [Type]
tps KernelBody (Aliases GPUMem)
body
    analyseExp (LastUse
lumap, Used
used) (Op (Inner (SegOp (SegHist lvl _ binops tps body)))) =
      let (LastUse
lumap', Used
used') = (HistOp (Aliases GPUMem) -> (LastUse, Used) -> (LastUse, Used))
-> (LastUse, Used) -> [HistOp (Aliases GPUMem)] -> (LastUse, Used)
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr HistOp (Aliases GPUMem) -> (LastUse, Used) -> (LastUse, Used)
analyseHistOp (LastUse
lumap, Used
used) [HistOp (Aliases GPUMem)]
binops
          (LastUse
lumap'', Used
used'') = (LastUse, Used) -> KernelBody (Aliases GPUMem) -> (LastUse, Used)
analyseKernelBody (LastUse
lumap', Used
used') KernelBody (Aliases GPUMem)
body
          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''
       in (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) (WithAcc [WithAccInput (Aliases GPUMem)]
_ Lambda (Aliases GPUMem)
l) =
      (LastUse, Used) -> Lambda (Aliases GPUMem) -> (LastUse, Used)
analyseLambda (LastUse
lumap, Used
used) Lambda (Aliases GPUMem)
l
    segOpHelper :: LastUse
-> Used
-> a
-> t (SegBinOp (Aliases GPUMem))
-> a
-> KernelBody (Aliases GPUMem)
-> (LastUse, Used)
segOpHelper LastUse
lumap Used
used a
lvl t (SegBinOp (Aliases GPUMem))
binops a
tps KernelBody (Aliases GPUMem)
body =
      let (LastUse
lumap', Used
used') = (SegBinOp (Aliases GPUMem) -> (LastUse, Used) -> (LastUse, Used))
-> (LastUse, Used)
-> t (SegBinOp (Aliases GPUMem))
-> (LastUse, Used)
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr SegBinOp (Aliases GPUMem) -> (LastUse, Used) -> (LastUse, Used)
analyseSegBinOp (LastUse
lumap, Used
used) t (SegBinOp (Aliases GPUMem))
binops
          (LastUse
lumap'', Used
used'') = (LastUse, Used) -> KernelBody (Aliases GPUMem) -> (LastUse, Used)
analyseKernelBody (LastUse
lumap', Used
used') KernelBody (Aliases GPUMem)
body
          nms :: Used
nms = (a -> Used
forall a. FreeIn a => a -> Used
freeIn a
lvl Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> a -> Used
forall a. FreeIn a => a -> Used
freeIn a
tps) Used -> Used -> Used
`namesSubtract` Used
used''
       in (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)

analyseBody :: LastUse -> Used -> Body (Aliases GPUMem) -> (LastUse, Used)
analyseBody :: LastUse -> Used -> Body (Aliases GPUMem) -> (LastUse, Used)
analyseBody LastUse
lumap Used
used (Body BodyDec (Aliases GPUMem)
_ Stms (Aliases GPUMem)
stms Result
result) =
  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
   in LastUse -> Used -> Stms (Aliases GPUMem) -> (LastUse, Used)
analyseStms LastUse
lumap Used
used' Stms (Aliases GPUMem)
stms

analyseKernelBody ::
  (LastUse, Used) ->
  KernelBody (Aliases GPUMem) ->
  (LastUse, Used)
analyseKernelBody :: (LastUse, Used) -> KernelBody (Aliases GPUMem) -> (LastUse, Used)
analyseKernelBody (LastUse
lumap, Used
used) (KernelBody BodyDec (Aliases GPUMem)
_ Stms (Aliases GPUMem)
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 GPUMem) -> (LastUse, Used)
analyseStms LastUse
lumap Used
used' Stms (Aliases GPUMem)
stms

analyseSegBinOp ::
  SegBinOp (Aliases GPUMem) ->
  (LastUse, Used) ->
  (LastUse, Used)
analyseSegBinOp :: SegBinOp (Aliases GPUMem) -> (LastUse, Used) -> (LastUse, Used)
analyseSegBinOp (SegBinOp Commutativity
_ Lambda (Aliases GPUMem)
lambda [SubExp]
neutral Shape
shp) (LastUse
lumap, Used
used) =
  let (LastUse
lumap', Used
used') = (LastUse, Used) -> Lambda (Aliases GPUMem) -> (LastUse, Used)
analyseLambda (LastUse
lumap, Used
used) Lambda (Aliases GPUMem)
lambda
      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'
   in (LastUse
lumap', Used
used' Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Used
nms)

analyseHistOp ::
  HistOp (Aliases GPUMem) ->
  (LastUse, Used) ->
  (LastUse, Used)
analyseHistOp :: HistOp (Aliases GPUMem) -> (LastUse, Used) -> (LastUse, Used)
analyseHistOp (HistOp SubExp
width SubExp
race [VName]
dest [SubExp]
neutral Shape
shp Lambda (Aliases GPUMem)
lambda) (LastUse
lumap, Used
used) =
  let (LastUse
lumap', Used
used') = (LastUse, Used) -> Lambda (Aliases GPUMem) -> (LastUse, Used)
analyseLambda (LastUse
lumap, Used
used) Lambda (Aliases GPUMem)
lambda
      nms :: Used
nms =
        ( SubExp -> Used
forall a. FreeIn a => a -> Used
freeIn SubExp
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'
   in (LastUse
lumap', Used
used' Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> Used
nms)

analyseLambda :: (LastUse, Used) -> Lambda (Aliases GPUMem) -> (LastUse, Used)
analyseLambda :: (LastUse, Used) -> Lambda (Aliases GPUMem) -> (LastUse, Used)
analyseLambda (LastUse
lumap, Used
used) (Lambda [LParam (Aliases GPUMem)]
params Body (Aliases GPUMem)
body [Type]
ret) =
  let (LastUse
lumap', Used
used') = LastUse -> Used -> Body (Aliases GPUMem) -> (LastUse, Used)
analyseBody LastUse
lumap Used
used Body (Aliases GPUMem)
body
      used'' :: Used
used'' = Used
used' Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> [Param LetDecMem] -> Used
forall a. FreeIn a => a -> Used
freeIn [LParam (Aliases GPUMem)]
[Param LetDecMem]
params Used -> Used -> Used
forall a. Semigroup a => a -> a -> a
<> [Type] -> Used
forall a. FreeIn a => a -> Used
freeIn [Type]
ret
   in (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