{-# 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 -> (Map VName Names, Names)
analyseProg Prog GPUMem
prog =
  let consts :: Names
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
. PatternT LetDecMem -> [PatElemT LetDecMem]
forall dec. PatternT dec -> [PatElemT dec]
patternValueElements (PatternT LetDecMem -> [PatElemT LetDecMem])
-> (Stm GPUMem -> PatternT LetDecMem)
-> Stm GPUMem
-> [PatElemT LetDecMem]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm GPUMem -> PatternT LetDecMem
forall rep. Stm rep -> Pattern rep
stmPattern)
          [VName] -> ([VName] -> Names) -> Names
forall a b. a -> (a -> b) -> b
& [VName] -> Names
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, Names
used) = (FunDef (Aliases GPUMem) -> (LastUse, Names))
-> [FunDef (Aliases GPUMem)] -> (LastUse, Names)
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (LastUse -> Names -> FunDef (Aliases GPUMem) -> (LastUse, Names)
analyseFun LastUse
forall a. Monoid a => a
mempty Names
consts) [FunDef (Aliases GPUMem)]
funs
   in (LastUse -> Map VName Names
flipMap LastUse
lus, Names
used)

analyseFun :: LastUse -> Used -> FunDef (Aliases GPUMem) -> (LastUse, Used)
analyseFun :: LastUse -> Names -> FunDef (Aliases GPUMem) -> (LastUse, Names)
analyseFun LastUse
lumap Names
used FunDef (Aliases GPUMem)
fun =
  let (LastUse
lumap', Names
used') = LastUse -> Names -> Body (Aliases GPUMem) -> (LastUse, Names)
analyseBody LastUse
lumap Names
used (Body (Aliases GPUMem) -> (LastUse, Names))
-> Body (Aliases GPUMem) -> (LastUse, Names)
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', Names
used' Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [Param FParamMem] -> Names
forall a. FreeIn a => a -> Names
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 -> Names -> Stms (Aliases GPUMem) -> (LastUse, Names)
analyseStms LastUse
lumap Names
used Stms (Aliases GPUMem)
stms = (Stm (Aliases GPUMem) -> (LastUse, Names) -> (LastUse, Names))
-> (LastUse, Names) -> [Stm (Aliases GPUMem)] -> (LastUse, Names)
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Stm (Aliases GPUMem) -> (LastUse, Names) -> (LastUse, Names)
analyseStm (LastUse
lumap, Names
used) ([Stm (Aliases GPUMem)] -> (LastUse, Names))
-> [Stm (Aliases GPUMem)] -> (LastUse, Names)
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, Names) -> (LastUse, Names)
analyseStm (Let Pattern (Aliases GPUMem)
pat StmAux (ExpDec (Aliases GPUMem))
_ Exp (Aliases GPUMem)
e) (LastUse
lumap0, Names
used0) =
  let (LastUse
lumap', Names
used') = PatternT (AliasDec, LetDecMem) -> [PatElemT (AliasDec, LetDecMem)]
forall dec. PatternT dec -> [PatElemT dec]
patternValueElements PatternT (AliasDec, LetDecMem)
Pattern (Aliases GPUMem)
pat [PatElemT (AliasDec, LetDecMem)]
-> ([PatElemT (AliasDec, LetDecMem)] -> (LastUse, Names))
-> (LastUse, Names)
forall a b. a -> (a -> b) -> b
& ((LastUse, Names)
 -> PatElemT (AliasDec, LetDecMem) -> (LastUse, Names))
-> (LastUse, Names)
-> [PatElemT (AliasDec, LetDecMem)]
-> (LastUse, Names)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (LastUse, Names)
-> PatElemT (AliasDec, LetDecMem) -> (LastUse, Names)
forall {b}.
(LastUse, Names) -> PatElemT (AliasDec, b) -> (LastUse, Names)
helper (LastUse
lumap0, Names
used0)
   in (LastUse, Names) -> Exp (Aliases GPUMem) -> (LastUse, Names)
analyseExp (LastUse
lumap', Names
used') Exp (Aliases GPUMem)
e
  where
    helper :: (LastUse, Names) -> PatElemT (AliasDec, b) -> (LastUse, Names)
helper (LastUse
lumap_acc, Names
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 -> Names -> LastUse -> LastUse
insertNames VName
name' (AliasDec -> Names
unAliases AliasDec
aliases) LastUse
lumap_acc
          Maybe VName
Nothing -> LastUse
lumap_acc,
        Names
used_acc Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> AliasDec -> Names
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
$ PatternT (AliasDec, LetDecMem) -> [PatElemT (AliasDec, LetDecMem)]
forall dec. PatternT dec -> [PatElemT dec]
patternValueElements PatternT (AliasDec, LetDecMem)
Pattern (Aliases GPUMem)
pat
    analyseExp :: (LastUse, Used) -> Exp (Aliases GPUMem) -> (LastUse, Used)
    analyseExp :: (LastUse, Names) -> Exp (Aliases GPUMem) -> (LastUse, Names)
analyseExp (LastUse
lumap, Names
used) (BasicOp BasicOp
_) =
      let nms :: Names
nms = Exp (Aliases GPUMem) -> Names
forall a. FreeIn a => a -> Names
freeIn Exp (Aliases GPUMem)
e Names -> Names -> Names
`namesSubtract` Names
used
       in (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 GPUMem)]
_ (Safety, SrcLoc, [SrcLoc])
_) =
      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
       in (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) (If SubExp
cse Body (Aliases GPUMem)
then_body Body (Aliases GPUMem)
else_body IfDec (BranchType (Aliases GPUMem))
dec) =
      let (LastUse
lumap_then, Names
used_then) = LastUse -> Names -> Body (Aliases GPUMem) -> (LastUse, Names)
analyseBody LastUse
lumap Names
used Body (Aliases GPUMem)
then_body
          (LastUse
lumap_else, Names
used_else) = LastUse -> Names -> Body (Aliases GPUMem) -> (LastUse, Names)
analyseBody LastUse
lumap Names
used Body (Aliases GPUMem)
else_body
          used' :: Names
used' = Names
used_then Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
used_else
          nms :: Names
nms = ((SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn SubExp
cse Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> IfDec BranchTypeMem -> Names
forall a. FreeIn a => a -> Names
freeIn IfDec (BranchType (Aliases GPUMem))
IfDec BranchTypeMem
dec) Names -> Names -> Names
`namesSubtract` Names
used')
       in (VName -> Names -> LastUse -> LastUse
insertNames VName
pat_name Names
nms (LastUse
lumap_then LastUse -> LastUse -> LastUse
forall a. Semigroup a => a -> a -> a
<> LastUse
lumap_else), Names
used' Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
nms)
    analyseExp (LastUse
lumap, Names
used) (DoLoop [(FParam (Aliases GPUMem), SubExp)]
ctx [(FParam (Aliases GPUMem), SubExp)]
vals LoopForm (Aliases GPUMem)
form Body (Aliases GPUMem)
body) =
      let (LastUse
lumap', Names
used') = LastUse -> Names -> Body (Aliases GPUMem) -> (LastUse, Names)
analyseBody LastUse
lumap Names
used Body (Aliases GPUMem)
body
          nms :: Names
nms = ([(Param FParamMem, SubExp)] -> Names
forall a. FreeIn a => a -> Names
freeIn [(FParam (Aliases GPUMem), SubExp)]
[(Param FParamMem, SubExp)]
ctx Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [(Param FParamMem, SubExp)] -> Names
forall a. FreeIn a => a -> Names
freeIn [(FParam (Aliases GPUMem), SubExp)]
[(Param FParamMem, SubExp)]
vals Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> LoopForm (Aliases GPUMem) -> Names
forall a. FreeIn a => a -> Names
freeIn LoopForm (Aliases GPUMem)
form) Names -> Names -> Names
`namesSubtract` Names
used'
       in (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 (Alloc SubExp
se Space
sp)) =
      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
       in (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 (Inner (SizeOp SizeOp
sop))) =
      let nms :: Names
nms = SizeOp -> Names
forall a. FreeIn a => a -> Names
freeIn SizeOp
sop Names -> Names -> Names
`namesSubtract` Names
used
       in (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 (Inner (OtherOp ()))) =
      (LastUse
lumap, Names
used)
    analyseExp (LastUse
lumap, Names
used) (Op (Inner (SegOp (SegMap SegLevel
lvl SegSpace
_ [Type]
tps KernelBody (Aliases GPUMem)
body)))) =
      let (LastUse
lumap', Names
used') = (LastUse, Names) -> KernelBody (Aliases GPUMem) -> (LastUse, Names)
analyseKernelBody (LastUse
lumap, Names
used) KernelBody (Aliases GPUMem)
body
          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'
       in (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 (Inner (SegOp (SegRed SegLevel
lvl SegSpace
_ [SegBinOp (Aliases GPUMem)]
binops [Type]
tps KernelBody (Aliases GPUMem)
body)))) =
      LastUse
-> Names
-> SegLevel
-> [SegBinOp (Aliases GPUMem)]
-> [Type]
-> KernelBody (Aliases GPUMem)
-> (LastUse, Names)
forall {t :: * -> *} {p} {p}.
(Foldable t, FreeIn p, FreeIn p) =>
LastUse
-> Names
-> p
-> t (SegBinOp (Aliases GPUMem))
-> p
-> KernelBody (Aliases GPUMem)
-> (LastUse, Names)
segOpHelper LastUse
lumap Names
used SegLevel
lvl [SegBinOp (Aliases GPUMem)]
binops [Type]
tps KernelBody (Aliases GPUMem)
body
    analyseExp (LastUse
lumap, Names
used) (Op (Inner (SegOp (SegScan SegLevel
lvl SegSpace
_ [SegBinOp (Aliases GPUMem)]
binops [Type]
tps KernelBody (Aliases GPUMem)
body)))) =
      LastUse
-> Names
-> SegLevel
-> [SegBinOp (Aliases GPUMem)]
-> [Type]
-> KernelBody (Aliases GPUMem)
-> (LastUse, Names)
forall {t :: * -> *} {p} {p}.
(Foldable t, FreeIn p, FreeIn p) =>
LastUse
-> Names
-> p
-> t (SegBinOp (Aliases GPUMem))
-> p
-> KernelBody (Aliases GPUMem)
-> (LastUse, Names)
segOpHelper LastUse
lumap Names
used SegLevel
lvl [SegBinOp (Aliases GPUMem)]
binops [Type]
tps KernelBody (Aliases GPUMem)
body
    analyseExp (LastUse
lumap, Names
used) (Op (Inner (SegOp (SegHist SegLevel
lvl SegSpace
_ [HistOp (Aliases GPUMem)]
binops [Type]
tps KernelBody (Aliases GPUMem)
body)))) =
      let (LastUse
lumap', Names
used') = (HistOp (Aliases GPUMem) -> (LastUse, Names) -> (LastUse, Names))
-> (LastUse, Names)
-> [HistOp (Aliases GPUMem)]
-> (LastUse, Names)
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr HistOp (Aliases GPUMem) -> (LastUse, Names) -> (LastUse, Names)
analyseHistOp (LastUse
lumap, Names
used) [HistOp (Aliases GPUMem)]
binops
          (LastUse
lumap'', Names
used'') = (LastUse, Names) -> KernelBody (Aliases GPUMem) -> (LastUse, Names)
analyseKernelBody (LastUse
lumap', Names
used') KernelBody (Aliases GPUMem)
body
          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''
       in (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) (WithAcc [(ShapeBase SubExp, [VName],
  Maybe (Lambda (Aliases GPUMem), [SubExp]))]
_ Lambda (Aliases GPUMem)
l) =
      (LastUse, Names) -> Lambda (Aliases GPUMem) -> (LastUse, Names)
analyseLambda (LastUse
lumap, Names
used) Lambda (Aliases GPUMem)
l
    segOpHelper :: LastUse
-> Names
-> p
-> t (SegBinOp (Aliases GPUMem))
-> p
-> KernelBody (Aliases GPUMem)
-> (LastUse, Names)
segOpHelper LastUse
lumap Names
used p
lvl t (SegBinOp (Aliases GPUMem))
binops p
tps KernelBody (Aliases GPUMem)
body =
      let (LastUse
lumap', Names
used') = (SegBinOp (Aliases GPUMem) -> (LastUse, Names) -> (LastUse, Names))
-> (LastUse, Names)
-> t (SegBinOp (Aliases GPUMem))
-> (LastUse, Names)
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr SegBinOp (Aliases GPUMem) -> (LastUse, Names) -> (LastUse, Names)
analyseSegBinOp (LastUse
lumap, Names
used) t (SegBinOp (Aliases GPUMem))
binops
          (LastUse
lumap'', Names
used'') = (LastUse, Names) -> KernelBody (Aliases GPUMem) -> (LastUse, Names)
analyseKernelBody (LastUse
lumap', Names
used') KernelBody (Aliases GPUMem)
body
          nms :: Names
nms = (p -> Names
forall a. FreeIn a => a -> Names
freeIn p
lvl Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> p -> Names
forall a. FreeIn a => a -> Names
freeIn p
tps) Names -> Names -> Names
`namesSubtract` Names
used''
       in (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)

analyseBody :: LastUse -> Used -> Body (Aliases GPUMem) -> (LastUse, Used)
analyseBody :: LastUse -> Names -> Body (Aliases GPUMem) -> (LastUse, Names)
analyseBody LastUse
lumap Names
used (Body BodyDec (Aliases GPUMem)
_ Stms (Aliases GPUMem)
stms [SubExp]
result) =
  let used' :: Names
used' = Names
used Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> Names
forall a. FreeIn a => a -> Names
freeIn [SubExp]
result
   in LastUse -> Names -> Stms (Aliases GPUMem) -> (LastUse, Names)
analyseStms LastUse
lumap Names
used' Stms (Aliases GPUMem)
stms

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

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

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

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