{-# LANGUAGE UndecidableInstances #-}

-- | This module implements common-subexpression elimination.  This
-- module does not actually remove the duplicate, but only replaces
-- one with a diference to the other.  E.g:
--
-- @
--   let a = x + y
--   let b = x + y
-- @
--
-- becomes:
--
-- @
--   let a = x + y
--   let b = a
-- @
--
-- After which copy propagation in the simplifier will actually remove
-- the definition of @b@.
--
-- Our CSE is still rather stupid.  No normalisation is performed, so
-- the expressions @x+y@ and @y+x@ will be considered distinct.
-- Furthermore, no expression with its own binding will be considered
-- equal to any other, since the variable names will be distinct.
-- This affects SOACs in particular.
module Futhark.Optimise.CSE
  ( performCSE,
    performCSEOnFunDef,
    performCSEOnStms,
    CSEInOp,
  )
where

import Control.Monad.Reader
import Data.Map.Strict qualified as M
import Futhark.Analysis.Alias
import Futhark.IR
import Futhark.IR.Aliases
  ( Aliases,
    consumedInStms,
    mkStmsAliases,
    removeFunDefAliases,
    removeProgAliases,
    removeStmAliases,
  )
import Futhark.IR.GPU qualified as GPU
import Futhark.IR.GPUMem qualified as GPUMem
import Futhark.IR.MC qualified as MC
import Futhark.IR.MCMem qualified as MCMem
import Futhark.IR.Mem qualified as Memory
import Futhark.IR.Prop.Aliases
import Futhark.IR.SOACS qualified as SOACS
import Futhark.IR.SOACS.SOAC qualified as SOAC
import Futhark.IR.Seq qualified as Seq
import Futhark.IR.SeqMem qualified as SeqMem
import Futhark.Pass
import Futhark.Transform.Substitute

-- | Perform CSE on every function in a program.
--
-- If the boolean argument is false, the pass will not perform CSE on
-- expressions producing arrays. This should be disabled when the rep has
-- memory information, since at that point arrays have identity beyond their
-- value.
performCSE ::
  (AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
  Bool ->
  Pass rep rep
performCSE :: forall rep.
(AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
Bool -> Pass rep rep
performCSE Bool
cse_arrays =
  String -> String -> (Prog rep -> PassM (Prog rep)) -> Pass rep rep
forall fromrep torep.
String
-> String
-> (Prog fromrep -> PassM (Prog torep))
-> Pass fromrep torep
Pass String
"CSE" String
"Combine common subexpressions." ((Prog rep -> PassM (Prog rep)) -> Pass rep rep)
-> (Prog rep -> PassM (Prog rep)) -> Pass rep rep
forall a b. (a -> b) -> a -> b
$ \Prog rep
prog ->
    (Prog (Aliases rep) -> Prog rep)
-> PassM (Prog (Aliases rep)) -> PassM (Prog rep)
forall a b. (a -> b) -> PassM a -> PassM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Prog (Aliases rep) -> Prog rep
forall rep. RephraseOp (OpC rep) => Prog (Aliases rep) -> Prog rep
removeProgAliases
      (PassM (Prog (Aliases rep)) -> PassM (Prog rep))
-> (Prog rep -> PassM (Prog (Aliases rep)))
-> Prog rep
-> PassM (Prog rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stms (Aliases rep) -> PassM (Stms (Aliases rep)))
-> (Stms (Aliases rep)
    -> FunDef (Aliases rep) -> PassM (FunDef (Aliases rep)))
-> Prog (Aliases rep)
-> PassM (Prog (Aliases rep))
forall fromrep torep.
(Stms fromrep -> PassM (Stms torep))
-> (Stms torep -> FunDef fromrep -> PassM (FunDef torep))
-> Prog fromrep
-> PassM (Prog torep)
intraproceduralTransformationWithConsts
        (Names -> Stms (Aliases rep) -> PassM (Stms (Aliases rep))
forall {f :: * -> *} {rep}.
(Applicative f, CSEInOp (OpC rep rep), Aliased rep) =>
Names -> Stms rep -> f (Stms rep)
onConsts ([FunDef rep] -> Names
forall a. FreeIn a => a -> Names
freeIn (Prog rep -> [FunDef rep]
forall rep. Prog rep -> [FunDef rep]
progFuns Prog rep
prog)))
        Stms (Aliases rep)
-> FunDef (Aliases rep) -> PassM (FunDef (Aliases rep))
forall {f :: * -> *} {rep} {p}.
(Applicative f, Aliased rep, CSEInOp (OpC rep rep)) =>
p -> FunDef rep -> f (FunDef rep)
onFun
      (Prog (Aliases rep) -> PassM (Prog (Aliases rep)))
-> (Prog rep -> Prog (Aliases rep))
-> Prog rep
-> PassM (Prog (Aliases rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog rep -> Prog (Aliases rep)
forall rep. AliasableRep rep => Prog rep -> Prog (Aliases rep)
aliasAnalysis
      (Prog rep -> PassM (Prog rep)) -> Prog rep -> PassM (Prog rep)
forall a b. (a -> b) -> a -> b
$ Prog rep
prog
  where
    onConsts :: Names -> Stms rep -> f (Stms rep)
onConsts Names
free_in_funs Stms rep
stms = do
      let free_list :: [VName]
free_list = Names -> [VName]
namesToList Names
free_in_funs
          ([Names]
res_als, Names
stms_cons) = Stms rep -> [SubExpRes] -> ([Names], Names)
forall rep res.
(Aliased rep, IsResult res) =>
Stms rep -> [res] -> ([Names], Names)
mkStmsAliases Stms rep
stms ([SubExpRes] -> ([Names], Names))
-> [SubExpRes] -> ([Names], Names)
forall a b. (a -> b) -> a -> b
$ [VName] -> [SubExpRes]
varsRes [VName]
free_list
      Stms rep -> f (Stms rep)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms rep -> f (Stms rep))
-> ((Stms rep, ()) -> Stms rep) -> (Stms rep, ()) -> f (Stms rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stms rep, ()) -> Stms rep
forall a b. (a, b) -> a
fst ((Stms rep, ()) -> f (Stms rep)) -> (Stms rep, ()) -> f (Stms rep)
forall a b. (a -> b) -> a -> b
$
        Reader (CSEState rep) (Stms rep, ())
-> CSEState rep -> (Stms rep, ())
forall r a. Reader r a -> r -> a
runReader
          ( Names
-> [Stm rep] -> CSEM rep () -> Reader (CSEState rep) (Stms rep, ())
forall rep a.
(Aliased rep, CSEInOp (Op rep)) =>
Names -> [Stm rep] -> CSEM rep a -> CSEM rep (Stms rep, a)
cseInStms
              ([Names] -> Names
forall a. Monoid a => [a] -> a
mconcat [Names]
res_als Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
stms_cons)
              (Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms rep
stms)
              (() -> CSEM rep ()
forall a. a -> ReaderT (CSEState rep) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
          )
          (Bool -> CSEState rep
forall rep. Bool -> CSEState rep
newCSEState Bool
cse_arrays)
    onFun :: p -> FunDef rep -> f (FunDef rep)
onFun p
_ = FunDef rep -> f (FunDef rep)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FunDef rep -> f (FunDef rep))
-> (FunDef rep -> FunDef rep) -> FunDef rep -> f (FunDef rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> FunDef rep -> FunDef rep
forall rep.
(Aliased rep, CSEInOp (Op rep)) =>
Bool -> FunDef rep -> FunDef rep
cseInFunDef Bool
cse_arrays

-- Specialising here means the importers that instantiate the pass (the
-- pipelines and the dev driver) use these copies instead of each
-- re-specialising the whole pass whenever they are recompiled.
{-# SPECIALIZE performCSE :: Bool -> Pass SOACS.SOACS SOACS.SOACS #-}
{-# SPECIALIZE performCSE :: Bool -> Pass GPU.GPU GPU.GPU #-}
{-# SPECIALIZE performCSE :: Bool -> Pass MC.MC MC.MC #-}
{-# SPECIALIZE performCSE :: Bool -> Pass Seq.Seq Seq.Seq #-}
{-# SPECIALIZE performCSE :: Bool -> Pass SeqMem.SeqMem SeqMem.SeqMem #-}
{-# SPECIALIZE performCSE :: Bool -> Pass GPUMem.GPUMem GPUMem.GPUMem #-}
{-# SPECIALIZE performCSE :: Bool -> Pass MCMem.MCMem MCMem.MCMem #-}

-- | Perform CSE on a single function.
--
-- If the boolean argument is false, the pass will not perform CSE on
-- expressions producing arrays. This should be disabled when the rep has
-- memory information, since at that point arrays have identity beyond their
-- value.
performCSEOnFunDef ::
  (AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
  Bool ->
  FunDef rep ->
  FunDef rep
performCSEOnFunDef :: forall rep.
(AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
Bool -> FunDef rep -> FunDef rep
performCSEOnFunDef Bool
cse_arrays =
  FunDef (Aliases rep) -> FunDef rep
forall rep.
RephraseOp (OpC rep) =>
FunDef (Aliases rep) -> FunDef rep
removeFunDefAliases (FunDef (Aliases rep) -> FunDef rep)
-> (FunDef rep -> FunDef (Aliases rep)) -> FunDef rep -> FunDef rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> FunDef (Aliases rep) -> FunDef (Aliases rep)
forall rep.
(Aliased rep, CSEInOp (Op rep)) =>
Bool -> FunDef rep -> FunDef rep
cseInFunDef Bool
cse_arrays (FunDef (Aliases rep) -> FunDef (Aliases rep))
-> (FunDef rep -> FunDef (Aliases rep))
-> FunDef rep
-> FunDef (Aliases rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FunDef rep -> FunDef (Aliases rep)
forall rep. AliasableRep rep => FunDef rep -> FunDef (Aliases rep)
analyseFun

-- | Perform CSE on some statements.
performCSEOnStms ::
  (AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
  Stms rep ->
  Stms rep
performCSEOnStms :: forall rep.
(AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
Stms rep -> Stms rep
performCSEOnStms =
  (Stm (Aliases rep) -> Stm rep)
-> Seq (Stm (Aliases rep)) -> Seq (Stm rep)
forall a b. (a -> b) -> Seq a -> Seq b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Stm (Aliases rep) -> Stm rep
forall rep. RephraseOp (OpC rep) => Stm (Aliases rep) -> Stm rep
removeStmAliases (Seq (Stm (Aliases rep)) -> Seq (Stm rep))
-> (Seq (Stm rep) -> Seq (Stm (Aliases rep)))
-> Seq (Stm rep)
-> Seq (Stm rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Seq (Stm (Aliases rep)) -> Seq (Stm (Aliases rep))
forall {rep}.
(Aliased rep, CSEInOp (OpC rep rep)) =>
Stms rep -> Stms rep
f (Seq (Stm (Aliases rep)) -> Seq (Stm (Aliases rep)))
-> (Seq (Stm rep) -> Seq (Stm (Aliases rep)))
-> Seq (Stm rep)
-> Seq (Stm (Aliases rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Seq (Stm (Aliases rep)), AliasesAndConsumed)
-> Seq (Stm (Aliases rep))
forall a b. (a, b) -> a
fst ((Seq (Stm (Aliases rep)), AliasesAndConsumed)
 -> Seq (Stm (Aliases rep)))
-> (Seq (Stm rep) -> (Seq (Stm (Aliases rep)), AliasesAndConsumed))
-> Seq (Stm rep)
-> Seq (Stm (Aliases rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AliasTable
-> Seq (Stm rep) -> (Seq (Stm (Aliases rep)), AliasesAndConsumed)
forall rep.
AliasableRep rep =>
AliasTable -> Stms rep -> (Stms (Aliases rep), AliasesAndConsumed)
analyseStms AliasTable
forall a. Monoid a => a
mempty
  where
    f :: Stms rep -> Stms rep
f Stms rep
stms =
      (Stms rep, ()) -> Stms rep
forall a b. (a, b) -> a
fst ((Stms rep, ()) -> Stms rep) -> (Stms rep, ()) -> Stms rep
forall a b. (a -> b) -> a -> b
$
        Reader (CSEState rep) (Stms rep, ())
-> CSEState rep -> (Stms rep, ())
forall r a. Reader r a -> r -> a
runReader
          (Names
-> [Stm rep] -> CSEM rep () -> Reader (CSEState rep) (Stms rep, ())
forall rep a.
(Aliased rep, CSEInOp (Op rep)) =>
Names -> [Stm rep] -> CSEM rep a -> CSEM rep (Stms rep, a)
cseInStms (Stms rep -> Names
forall rep. Aliased rep => Stms rep -> Names
consumedInStms Stms rep
stms) (Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms rep
stms) (() -> CSEM rep ()
forall a. a -> ReaderT (CSEState rep) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()))
          -- It is never safe to CSE arrays in stms in isolation,
          -- because we might introduce additional aliasing.
          (Bool -> CSEState rep
forall rep. Bool -> CSEState rep
newCSEState Bool
False)

cseInFunDef ::
  (Aliased rep, CSEInOp (Op rep)) =>
  Bool ->
  FunDef rep ->
  FunDef rep
cseInFunDef :: forall rep.
(Aliased rep, CSEInOp (Op rep)) =>
Bool -> FunDef rep -> FunDef rep
cseInFunDef Bool
cse_arrays FunDef rep
fundec =
  FunDef rep
fundec
    { funDefBody =
        runReader (cseInBody ds $ funDefBody fundec) $ newCSEState cse_arrays
    }
  where
    -- XXX: we treat every array result as a consumption here, because
    -- it is otherwise complicated to ensure we do not introduce more
    -- aliasing than specified by the return type. This is not a
    -- practical problem while we still perform such aggressive
    -- inlining.
    ds :: [Diet]
ds = ((RetType rep, RetAls) -> Diet)
-> [(RetType rep, RetAls)] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map (RetType rep -> Diet
forall {t}. ExtTyped t => t -> Diet
retDiet (RetType rep -> Diet)
-> ((RetType rep, RetAls) -> RetType rep)
-> (RetType rep, RetAls)
-> Diet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (RetType rep, RetAls) -> RetType rep
forall a b. (a, b) -> a
fst) ([(RetType rep, RetAls)] -> [Diet])
-> [(RetType rep, RetAls)] -> [Diet]
forall a b. (a -> b) -> a -> b
$ FunDef rep -> [(RetType rep, RetAls)]
forall rep. FunDef rep -> [(RetType rep, RetAls)]
funDefRetType FunDef rep
fundec
    retDiet :: t -> Diet
retDiet t
t
      | TypeBase ExtShape NoMode -> Bool
forall shape o. TypeBase shape o -> Bool
primType (TypeBase ExtShape NoMode -> Bool)
-> TypeBase ExtShape NoMode -> Bool
forall a b. (a -> b) -> a -> b
$ t -> TypeBase ExtShape NoMode
forall t. ExtTyped t => t -> TypeBase ExtShape NoMode
extTypeOf t
t = Diet
Observe
      | Bool
otherwise = Diet
Consume

type CSEM rep = Reader (CSEState rep)

-- | Like 'substituteNames', but leaves the term alone rather than
-- rebuilding an identical one when there is nothing to substitute.
substituteNames' :: (Substitute a) => M.Map VName VName -> a -> a
substituteNames' :: forall a. Substitute a => Map VName VName -> a -> a
substituteNames' Map VName VName
substs a
x
  | Map VName VName -> Bool
forall k a. Map k a -> Bool
M.null Map VName VName
substs = a
x
  | Bool
otherwise = Map VName VName -> a -> a
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs a
x

cseInBody ::
  (Aliased rep, CSEInOp (Op rep)) =>
  [Diet] ->
  Body rep ->
  CSEM rep (Body rep)
cseInBody :: forall rep.
(Aliased rep, CSEInOp (Op rep)) =>
[Diet] -> Body rep -> CSEM rep (Body rep)
cseInBody [Diet]
ds (Body BodyDec rep
bodydec Stms rep
stms [SubExpRes]
res) = do
  (stms', res') <-
    Names
-> [Stm rep]
-> CSEM rep [SubExpRes]
-> CSEM rep (Stms rep, [SubExpRes])
forall rep a.
(Aliased rep, CSEInOp (Op rep)) =>
Names -> [Stm rep] -> CSEM rep a -> CSEM rep (Stms rep, a)
cseInStms (Names
res_cons Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
stms_cons) (Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms rep
stms) (CSEM rep [SubExpRes] -> CSEM rep (Stms rep, [SubExpRes]))
-> CSEM rep [SubExpRes] -> CSEM rep (Stms rep, [SubExpRes])
forall a b. (a -> b) -> a -> b
$ do
      CSEState (_, nsubsts) _ <- ReaderT (CSEState rep) Identity (CSEState rep)
forall r (m :: * -> *). MonadReader r m => m r
ask
      pure $ substituteNames' nsubsts res
  pure $ Body bodydec stms' res'
  where
    ([Names]
res_als, Names
stms_cons) = Stms rep -> [SubExpRes] -> ([Names], Names)
forall rep res.
(Aliased rep, IsResult res) =>
Stms rep -> [res] -> ([Names], Names)
mkStmsAliases Stms rep
stms [SubExpRes]
res
    res_cons :: Names
res_cons = [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat ([Names] -> Names) -> [Names] -> Names
forall a b. (a -> b) -> a -> b
$ (Diet -> Names -> Names) -> [Diet] -> [Names] -> [Names]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Diet -> Names -> Names
forall {p}. Monoid p => Diet -> p -> p
consumeResult [Diet]
ds [Names]
res_als
    consumeResult :: Diet -> p -> p
consumeResult Diet
Consume p
als = p
als
    consumeResult Diet
_ p
_ = p
forall a. Monoid a => a
mempty

cseInLambda ::
  (Aliased rep, CSEInOp (Op rep)) =>
  GLambda rep t ->
  CSEM rep (GLambda rep t)
cseInLambda :: forall rep t.
(Aliased rep, CSEInOp (Op rep)) =>
GLambda rep t -> CSEM rep (GLambda rep t)
cseInLambda GLambda rep t
lam = do
  body' <- [Diet] -> Body rep -> CSEM rep (Body rep)
forall rep.
(Aliased rep, CSEInOp (Op rep)) =>
[Diet] -> Body rep -> CSEM rep (Body rep)
cseInBody ((t -> Diet) -> [t] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map (Diet -> t -> Diet
forall a b. a -> b -> a
const Diet
Observe) ([t] -> [Diet]) -> [t] -> [Diet]
forall a b. (a -> b) -> a -> b
$ GLambda rep t -> [t]
forall rep t. GLambda rep t -> [t]
lambdaReturnType GLambda rep t
lam) (Body rep -> CSEM rep (Body rep))
-> Body rep -> CSEM rep (Body rep)
forall a b. (a -> b) -> a -> b
$ GLambda rep t -> Body rep
forall rep t. GLambda rep t -> Body rep
lambdaBody GLambda rep t
lam
  pure lam {lambdaBody = body'}

cseInStms ::
  forall rep a.
  (Aliased rep, CSEInOp (Op rep)) =>
  Names ->
  [Stm rep] ->
  CSEM rep a ->
  CSEM rep (Stms rep, a)
cseInStms :: forall rep a.
(Aliased rep, CSEInOp (Op rep)) =>
Names -> [Stm rep] -> CSEM rep a -> CSEM rep (Stms rep, a)
cseInStms Names
_ [] CSEM rep a
m = do
  a <- CSEM rep a
m
  pure (mempty, a)
cseInStms Names
consumed (Stm rep
stm : [Stm rep]
stms) CSEM rep a
m =
  Names
-> Stm rep
-> ([Stm rep] -> ReaderT (CSEState rep) Identity (Stms rep, a))
-> ReaderT (CSEState rep) Identity (Stms rep, a)
forall rep a.
ASTRep rep =>
Names -> Stm rep -> ([Stm rep] -> CSEM rep a) -> CSEM rep a
cseInStm Names
consumed Stm rep
stm (([Stm rep] -> ReaderT (CSEState rep) Identity (Stms rep, a))
 -> ReaderT (CSEState rep) Identity (Stms rep, a))
-> ([Stm rep] -> ReaderT (CSEState rep) Identity (Stms rep, a))
-> ReaderT (CSEState rep) Identity (Stms rep, a)
forall a b. (a -> b) -> a -> b
$ \[Stm rep]
stm' -> do
    (stms', a) <- Names
-> [Stm rep]
-> CSEM rep a
-> ReaderT (CSEState rep) Identity (Stms rep, a)
forall rep a.
(Aliased rep, CSEInOp (Op rep)) =>
Names -> [Stm rep] -> CSEM rep a -> CSEM rep (Stms rep, a)
cseInStms Names
consumed [Stm rep]
stms CSEM rep a
m
    stm'' <- mapM nestedCSE stm'
    pure (stmsFromList stm'' <> stms', a)
  where
    nestedCSE :: Stm rep -> ReaderT (CSEState rep) Identity (Stm rep)
nestedCSE Stm rep
stm' = do
      let ds :: [Diet]
ds =
            case Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm' of
              Loop [(Param (FParamInfo rep), SubExp)]
merge LoopForm
_ Body rep
_ -> ((Param (FParamInfo rep), SubExp) -> Diet)
-> [(Param (FParamInfo rep), SubExp)] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase Shape Diet -> Diet
forall shape. TypeBase shape Diet -> Diet
diet (TypeBase Shape Diet -> Diet)
-> ((Param (FParamInfo rep), SubExp) -> TypeBase Shape Diet)
-> (Param (FParamInfo rep), SubExp)
-> Diet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param (FParamInfo rep) -> TypeBase Shape Diet
forall t. DeclTyped t => t -> TypeBase Shape Diet
declTypeOf (Param (FParamInfo rep) -> TypeBase Shape Diet)
-> ((Param (FParamInfo rep), SubExp) -> Param (FParamInfo rep))
-> (Param (FParamInfo rep), SubExp)
-> TypeBase Shape Diet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param (FParamInfo rep), SubExp) -> Param (FParamInfo rep)
forall a b. (a, b) -> a
fst) [(Param (FParamInfo rep), SubExp)]
merge
              Exp rep
_ -> (PatElem (LetDec rep) -> Diet) -> [PatElem (LetDec rep)] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map PatElem (LetDec rep) -> Diet
forall {dec}. PatElem dec -> Diet
patElemDiet ([PatElem (LetDec rep)] -> [Diet])
-> [PatElem (LetDec rep)] -> [Diet]
forall a b. (a -> b) -> a -> b
$ Pat (LetDec rep) -> [PatElem (LetDec rep)]
forall dec. Pat dec -> [PatElem dec]
patElems (Pat (LetDec rep) -> [PatElem (LetDec rep)])
-> Pat (LetDec rep) -> [PatElem (LetDec rep)]
forall a b. (a -> b) -> a -> b
$ Stm rep -> Pat (LetDec rep)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm rep
stm'
      e <- Mapper rep rep (ReaderT (CSEState rep) Identity)
-> Exp rep -> ReaderT (CSEState rep) Identity (Exp rep)
forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM ([Diet] -> Mapper rep rep (ReaderT (CSEState rep) Identity)
forall {frep}.
(FParamInfo frep ~ FParamInfo rep, RetType frep ~ RetType rep,
 BranchType frep ~ BranchType rep, LParamInfo frep ~ LParamInfo rep,
 Aliased frep, CSEInOp (OpC frep frep)) =>
[Diet] -> Mapper frep frep (ReaderT (CSEState frep) Identity)
cse [Diet]
ds) (Exp rep -> ReaderT (CSEState rep) Identity (Exp rep))
-> Exp rep -> ReaderT (CSEState rep) Identity (Exp rep)
forall a b. (a -> b) -> a -> b
$ Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm'
      pure stm' {stmExp = e}

    cse :: [Diet] -> Mapper frep frep (ReaderT (CSEState frep) Identity)
cse [Diet]
ds =
      (forall rep (m :: * -> *). Monad m => Mapper rep rep m
identityMapper @rep)
        { mapOnBody = const $ cseInBody ds,
          mapOnOp = cseInOp
        }

    patElemDiet :: PatElem dec -> Diet
patElemDiet PatElem dec
pe
      | PatElem dec -> VName
forall dec. PatElem dec -> VName
patElemName PatElem dec
pe VName -> Names -> Bool
`nameIn` Names
consumed = Diet
Consume
      | Bool
otherwise = Diet
Observe

cseInStm ::
  (ASTRep rep) =>
  Names ->
  Stm rep ->
  ([Stm rep] -> CSEM rep a) ->
  CSEM rep a
cseInStm :: forall rep a.
ASTRep rep =>
Names -> Stm rep -> ([Stm rep] -> CSEM rep a) -> CSEM rep a
cseInStm Names
consumed (Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
aux Exp rep
e) [Stm rep] -> CSEM rep a
m = do
  CSEState (esubsts, nsubsts) cse_arrays <- ReaderT (CSEState rep) Identity (CSEState rep)
forall r (m :: * -> *). MonadReader r m => m r
ask
  let e' = Map VName VName -> Exp rep -> Exp rep
forall a. Substitute a => Map VName VName -> a -> a
substituteNames' Map VName VName
nsubsts Exp rep
e
      pat' = Map VName VName -> Pat (LetDec rep) -> Pat (LetDec rep)
forall a. Substitute a => Map VName VName -> a -> a
substituteNames' Map VName VName
nsubsts Pat (LetDec rep)
pat
  if not (alreadyAliases e) && any (bad cse_arrays) (patElems pat)
    then m [Let pat' aux e']
    else case M.lookup (stmAuxDec aux, e') esubsts of
      Just (Certs
subcs, Pat (LetDec rep)
subpat) -> do
        let subsumes :: Bool
subsumes = (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` Certs -> [VName]
unCerts Certs
subcs) (Certs -> [VName]
unCerts (StmAux (ExpDec rep) -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts StmAux (ExpDec rep)
aux))
        -- We can only do a plain name substitution if it doesn't
        -- violate any certificate dependencies.
        (CSEState rep -> CSEState rep) -> CSEM rep a -> CSEM rep a
forall a.
(CSEState rep -> CSEState rep)
-> ReaderT (CSEState rep) Identity a
-> ReaderT (CSEState rep) Identity a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (if Bool
subsumes then Pat (LetDec rep)
-> Pat (LetDec rep) -> CSEState rep -> CSEState rep
forall dec rep. Pat dec -> Pat dec -> CSEState rep -> CSEState rep
addNameSubst Pat (LetDec rep)
pat' Pat (LetDec rep)
subpat else CSEState rep -> CSEState rep
forall a. a -> a
id) (CSEM rep a -> CSEM rep a) -> CSEM rep a -> CSEM rep a
forall a b. (a -> b) -> a -> b
$ do
          let lets :: [Stm rep]
lets =
                [ Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let ([PatElem (LetDec rep)] -> Pat (LetDec rep)
forall dec. [PatElem dec] -> Pat dec
Pat [PatElem (LetDec rep)
patElem']) StmAux (ExpDec rep)
aux (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$
                    BasicOp -> Exp rep
forall rep. BasicOp -> Exp rep
BasicOp (SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (LetDec rep)
patElem)
                | (VName
name, PatElem (LetDec rep)
patElem) <- [VName]
-> [PatElem (LetDec rep)] -> [(VName, PatElem (LetDec rep))]
forall a b. [a] -> [b] -> [(a, b)]
zip (Pat (LetDec rep) -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat (LetDec rep)
pat') ([PatElem (LetDec rep)] -> [(VName, PatElem (LetDec rep))])
-> [PatElem (LetDec rep)] -> [(VName, PatElem (LetDec rep))]
forall a b. (a -> b) -> a -> b
$ Pat (LetDec rep) -> [PatElem (LetDec rep)]
forall dec. Pat dec -> [PatElem dec]
patElems Pat (LetDec rep)
subpat,
                  let patElem' :: PatElem (LetDec rep)
patElem' = PatElem (LetDec rep)
patElem {patElemName = name}
                ]
          [Stm rep] -> CSEM rep a
m [Stm rep]
lets
      Maybe (Certs, Pat (LetDec rep))
_ ->
        (CSEState rep -> CSEState rep) -> CSEM rep a -> CSEM rep a
forall a.
(CSEState rep -> CSEState rep)
-> ReaderT (CSEState rep) Identity a
-> ReaderT (CSEState rep) Identity a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (Pat (LetDec rep)
-> ExpDec rep -> Certs -> Exp rep -> CSEState rep -> CSEState rep
forall rep.
ASTRep rep =>
Pat (LetDec rep)
-> ExpDec rep -> Certs -> Exp rep -> CSEState rep -> CSEState rep
addExpSubst Pat (LetDec rep)
pat' (StmAux (ExpDec rep) -> ExpDec rep
forall dec. StmAux dec -> dec
stmAuxDec StmAux (ExpDec rep)
aux) (StmAux (ExpDec rep) -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts StmAux (ExpDec rep)
aux) Exp rep
e') (CSEM rep a -> CSEM rep a) -> CSEM rep a -> CSEM rep a
forall a b. (a -> b) -> a -> b
$
          [Stm rep] -> CSEM rep a
m [Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec rep)
pat' StmAux (ExpDec rep)
aux Exp rep
e']
  where
    alreadyAliases :: Exp rep -> Bool
alreadyAliases (BasicOp Index {}) = Bool
True
    alreadyAliases (BasicOp Reshape {}) = Bool
True
    alreadyAliases Exp rep
_ = Bool
False
    bad :: Bool -> PatElem dec -> Bool
bad Bool
cse_arrays PatElem dec
pe
      | Mem {} <- PatElem dec -> TypeBase Shape NoMode
forall dec. Typed dec => PatElem dec -> TypeBase Shape NoMode
patElemType PatElem dec
pe = Bool
True
      | Array {} <- PatElem dec -> TypeBase Shape NoMode
forall dec. Typed dec => PatElem dec -> TypeBase Shape NoMode
patElemType PatElem dec
pe, Bool -> Bool
not Bool
cse_arrays = Bool
True
      | PatElem dec -> VName
forall dec. PatElem dec -> VName
patElemName PatElem dec
pe VName -> Names -> Bool
`nameIn` Names
consumed = Bool
True
      | Bool
otherwise = Bool
False

type ExpressionSubstitutions rep =
  M.Map
    (ExpDec rep, Exp rep)
    (Certs, Pat (LetDec rep))

type NameSubstitutions = M.Map VName VName

data CSEState rep = CSEState
  { forall rep.
CSEState rep -> (ExpressionSubstitutions rep, Map VName VName)
_cseSubstitutions :: (ExpressionSubstitutions rep, NameSubstitutions),
    forall rep. CSEState rep -> Bool
_cseArrays :: Bool
  }

newCSEState :: Bool -> CSEState rep
newCSEState :: forall rep. Bool -> CSEState rep
newCSEState = (ExpressionSubstitutions rep, Map VName VName)
-> Bool -> CSEState rep
forall rep.
(ExpressionSubstitutions rep, Map VName VName)
-> Bool -> CSEState rep
CSEState (ExpressionSubstitutions rep
forall k a. Map k a
M.empty, Map VName VName
forall k a. Map k a
M.empty)

mkSubsts :: Pat dec -> Pat dec -> M.Map VName VName
mkSubsts :: forall dec. Pat dec -> Pat dec -> Map VName VName
mkSubsts Pat dec
pat Pat dec
vs = [(VName, VName)] -> Map VName VName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, VName)] -> Map VName VName)
-> [(VName, VName)] -> Map VName VName
forall a b. (a -> b) -> a -> b
$ [VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Pat dec -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat dec
pat) (Pat dec -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat dec
vs)

addNameSubst :: Pat dec -> Pat dec -> CSEState rep -> CSEState rep
addNameSubst :: forall dec rep. Pat dec -> Pat dec -> CSEState rep -> CSEState rep
addNameSubst Pat dec
pat Pat dec
subpat (CSEState (ExpressionSubstitutions rep
esubsts, Map VName VName
nsubsts) Bool
cse_arrays) =
  (ExpressionSubstitutions rep, Map VName VName)
-> Bool -> CSEState rep
forall rep.
(ExpressionSubstitutions rep, Map VName VName)
-> Bool -> CSEState rep
CSEState (ExpressionSubstitutions rep
esubsts, Pat dec -> Pat dec -> Map VName VName
forall dec. Pat dec -> Pat dec -> Map VName VName
mkSubsts Pat dec
pat Pat dec
subpat Map VName VName -> Map VName VName -> Map VName VName
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` Map VName VName
nsubsts) Bool
cse_arrays

addExpSubst ::
  (ASTRep rep) =>
  Pat (LetDec rep) ->
  ExpDec rep ->
  Certs ->
  Exp rep ->
  CSEState rep ->
  CSEState rep
addExpSubst :: forall rep.
ASTRep rep =>
Pat (LetDec rep)
-> ExpDec rep -> Certs -> Exp rep -> CSEState rep -> CSEState rep
addExpSubst Pat (LetDec rep)
pat ExpDec rep
edec Certs
cs Exp rep
e (CSEState (ExpressionSubstitutions rep
esubsts, Map VName VName
nsubsts) Bool
cse_arrays) =
  (ExpressionSubstitutions rep, Map VName VName)
-> Bool -> CSEState rep
forall rep.
(ExpressionSubstitutions rep, Map VName VName)
-> Bool -> CSEState rep
CSEState ((ExpDec rep, Exp rep)
-> (Certs, Pat (LetDec rep))
-> ExpressionSubstitutions rep
-> ExpressionSubstitutions rep
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (ExpDec rep
edec, Exp rep
e) (Certs
cs, Pat (LetDec rep)
pat) ExpressionSubstitutions rep
esubsts, Map VName VName
nsubsts) Bool
cse_arrays

-- | The operations that permit CSE.
class CSEInOp op where
  -- | Perform CSE within any nested expressions.
  cseInOp :: op -> CSEM rep op

instance CSEInOp (NoOp rep) where
  cseInOp :: forall rep. NoOp rep -> CSEM rep (NoOp rep)
cseInOp NoOp rep
NoOp = NoOp rep -> ReaderT (CSEState rep) Identity (NoOp rep)
forall a. a -> ReaderT (CSEState rep) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NoOp rep
forall {k} (rep :: k). NoOp rep
NoOp

subCSE :: CSEM rep r -> CSEM otherrep r
subCSE :: forall rep r otherrep. CSEM rep r -> CSEM otherrep r
subCSE CSEM rep r
m = do
  CSEState _ cse_arrays <- ReaderT (CSEState otherrep) Identity (CSEState otherrep)
forall r (m :: * -> *). MonadReader r m => m r
ask
  pure $ runReader m $ newCSEState cse_arrays

instance
  ( Aliased rep,
    CSEInOp (Op rep),
    CSEInOp (op rep)
  ) =>
  CSEInOp (GPU.HostOp op rep)
  where
  cseInOp :: forall rep. HostOp op rep -> CSEM rep (HostOp op rep)
cseInOp (GPU.SegOp SegOp SegLevel rep
op) = SegOp SegLevel rep -> HostOp op rep
forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
GPU.SegOp (SegOp SegLevel rep -> HostOp op rep)
-> ReaderT (CSEState rep) Identity (SegOp SegLevel rep)
-> ReaderT (CSEState rep) Identity (HostOp op rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegOp SegLevel rep
-> ReaderT (CSEState rep) Identity (SegOp SegLevel rep)
forall rep. SegOp SegLevel rep -> CSEM rep (SegOp SegLevel rep)
forall op rep. CSEInOp op => op -> CSEM rep op
cseInOp SegOp SegLevel rep
op
  cseInOp (GPU.OtherOp op rep
op) = op rep -> HostOp op rep
forall (op :: * -> *) rep. op rep -> HostOp op rep
GPU.OtherOp (op rep -> HostOp op rep)
-> ReaderT (CSEState rep) Identity (op rep)
-> ReaderT (CSEState rep) Identity (HostOp op rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> op rep -> ReaderT (CSEState rep) Identity (op rep)
forall rep. op rep -> CSEM rep (op rep)
forall op rep. CSEInOp op => op -> CSEM rep op
cseInOp op rep
op
  cseInOp (GPU.GPUBody [TypeBase Shape NoMode]
types Body rep
body) =
    CSEM rep (HostOp op rep)
-> ReaderT (CSEState rep) Identity (HostOp op rep)
forall rep r otherrep. CSEM rep r -> CSEM otherrep r
subCSE (CSEM rep (HostOp op rep)
 -> ReaderT (CSEState rep) Identity (HostOp op rep))
-> CSEM rep (HostOp op rep)
-> ReaderT (CSEState rep) Identity (HostOp op rep)
forall a b. (a -> b) -> a -> b
$ [TypeBase Shape NoMode] -> Body rep -> HostOp op rep
forall (op :: * -> *) rep.
[TypeBase Shape NoMode] -> Body rep -> HostOp op rep
GPU.GPUBody [TypeBase Shape NoMode]
types (Body rep -> HostOp op rep)
-> ReaderT (CSEState rep) Identity (Body rep)
-> CSEM rep (HostOp op rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Diet] -> Body rep -> ReaderT (CSEState rep) Identity (Body rep)
forall rep.
(Aliased rep, CSEInOp (Op rep)) =>
[Diet] -> Body rep -> CSEM rep (Body rep)
cseInBody ((TypeBase Shape NoMode -> Diet)
-> [TypeBase Shape NoMode] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map (Diet -> TypeBase Shape NoMode -> Diet
forall a b. a -> b -> a
const Diet
Observe) [TypeBase Shape NoMode]
types) Body rep
body
  cseInOp HostOp op rep
x = HostOp op rep -> ReaderT (CSEState rep) Identity (HostOp op rep)
forall a. a -> ReaderT (CSEState rep) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure HostOp op rep
x

instance
  ( Aliased rep,
    CSEInOp (Op rep),
    CSEInOp (op rep)
  ) =>
  CSEInOp (MC.MCOp op rep)
  where
  cseInOp :: forall rep. MCOp op rep -> CSEM rep (MCOp op rep)
cseInOp (MC.ParOp Maybe (SegOp () rep)
par_op SegOp () rep
op) =
    Maybe (SegOp () rep) -> SegOp () rep -> MCOp op rep
forall (op :: * -> *) rep.
Maybe (SegOp () rep) -> SegOp () rep -> MCOp op rep
MC.ParOp (Maybe (SegOp () rep) -> SegOp () rep -> MCOp op rep)
-> ReaderT (CSEState rep) Identity (Maybe (SegOp () rep))
-> ReaderT (CSEState rep) Identity (SegOp () rep -> MCOp op rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SegOp () rep -> ReaderT (CSEState rep) Identity (SegOp () rep))
-> Maybe (SegOp () rep)
-> ReaderT (CSEState rep) Identity (Maybe (SegOp () rep))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse SegOp () rep -> ReaderT (CSEState rep) Identity (SegOp () rep)
forall rep. SegOp () rep -> CSEM rep (SegOp () rep)
forall op rep. CSEInOp op => op -> CSEM rep op
cseInOp Maybe (SegOp () rep)
par_op ReaderT (CSEState rep) Identity (SegOp () rep -> MCOp op rep)
-> ReaderT (CSEState rep) Identity (SegOp () rep)
-> ReaderT (CSEState rep) Identity (MCOp op rep)
forall a b.
ReaderT (CSEState rep) Identity (a -> b)
-> ReaderT (CSEState rep) Identity a
-> ReaderT (CSEState rep) Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SegOp () rep -> ReaderT (CSEState rep) Identity (SegOp () rep)
forall rep. SegOp () rep -> CSEM rep (SegOp () rep)
forall op rep. CSEInOp op => op -> CSEM rep op
cseInOp SegOp () rep
op
  cseInOp (MC.OtherOp op rep
op) =
    op rep -> MCOp op rep
forall (op :: * -> *) rep. op rep -> MCOp op rep
MC.OtherOp (op rep -> MCOp op rep)
-> ReaderT (CSEState rep) Identity (op rep)
-> ReaderT (CSEState rep) Identity (MCOp op rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> op rep -> ReaderT (CSEState rep) Identity (op rep)
forall rep. op rep -> CSEM rep (op rep)
forall op rep. CSEInOp op => op -> CSEM rep op
cseInOp op rep
op

instance
  (Aliased rep, CSEInOp (Op rep)) =>
  CSEInOp (GPU.SegOp lvl rep)
  where
  cseInOp :: forall rep. SegOp lvl rep -> CSEM rep (SegOp lvl rep)
cseInOp =
    CSEM rep (SegOp lvl rep) -> CSEM rep (SegOp lvl rep)
forall rep r otherrep. CSEM rep r -> CSEM otherrep r
subCSE
      (CSEM rep (SegOp lvl rep) -> CSEM rep (SegOp lvl rep))
-> (SegOp lvl rep -> CSEM rep (SegOp lvl rep))
-> SegOp lvl rep
-> CSEM rep (SegOp lvl rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SegOpMapper lvl rep rep (ReaderT (CSEState rep) Identity)
-> SegOp lvl rep -> CSEM rep (SegOp lvl rep)
forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
GPU.mapSegOpM
        ((SubExp -> ReaderT (CSEState rep) Identity SubExp)
-> (Lambda rep -> ReaderT (CSEState rep) Identity (Lambda rep))
-> (Lambda rep -> ReaderT (CSEState rep) Identity (Lambda rep))
-> (KernelBody rep
    -> ReaderT (CSEState rep) Identity (KernelBody rep))
-> (VName -> ReaderT (CSEState rep) Identity VName)
-> (lvl -> ReaderT (CSEState rep) Identity lvl)
-> SegOpMapper lvl rep rep (ReaderT (CSEState rep) Identity)
forall lvl frep trep (m :: * -> *).
(SubExp -> m SubExp)
-> (Lambda frep -> m (Lambda trep))
-> (Lambda frep -> m (Lambda trep))
-> (KernelBody frep -> m (KernelBody trep))
-> (VName -> m VName)
-> (lvl -> m lvl)
-> SegOpMapper lvl frep trep m
GPU.SegOpMapper SubExp -> ReaderT (CSEState rep) Identity SubExp
forall a. a -> ReaderT (CSEState rep) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda rep -> ReaderT (CSEState rep) Identity (Lambda rep)
forall rep t.
(Aliased rep, CSEInOp (Op rep)) =>
GLambda rep t -> CSEM rep (GLambda rep t)
cseInLambda Lambda rep -> ReaderT (CSEState rep) Identity (Lambda rep)
forall rep t.
(Aliased rep, CSEInOp (Op rep)) =>
GLambda rep t -> CSEM rep (GLambda rep t)
cseInLambda KernelBody rep -> ReaderT (CSEState rep) Identity (KernelBody rep)
forall rep.
(Aliased rep, CSEInOp (Op rep)) =>
KernelBody rep -> CSEM rep (KernelBody rep)
cseInKernelBody VName -> ReaderT (CSEState rep) Identity VName
forall a. a -> ReaderT (CSEState rep) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure lvl -> ReaderT (CSEState rep) Identity lvl
forall a. a -> ReaderT (CSEState rep) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure)

cseInKernelBody ::
  (Aliased rep, CSEInOp (Op rep)) =>
  GPU.KernelBody rep ->
  CSEM rep (GPU.KernelBody rep)
cseInKernelBody :: forall rep.
(Aliased rep, CSEInOp (Op rep)) =>
KernelBody rep -> CSEM rep (KernelBody rep)
cseInKernelBody (GPU.Body BodyDec rep
bodydec Stms rep
stms [KernelResult]
res) = do
  Body _ stms' _ <- [Diet] -> GBody rep SubExpRes -> CSEM rep (GBody rep SubExpRes)
forall rep.
(Aliased rep, CSEInOp (Op rep)) =>
[Diet] -> Body rep -> CSEM rep (Body rep)
cseInBody ((KernelResult -> Diet) -> [KernelResult] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map (Diet -> KernelResult -> Diet
forall a b. a -> b -> a
const Diet
Observe) [KernelResult]
res) (GBody rep SubExpRes -> CSEM rep (GBody rep SubExpRes))
-> GBody rep SubExpRes -> CSEM rep (GBody rep SubExpRes)
forall a b. (a -> b) -> a -> b
$ BodyDec rep -> Stms rep -> [SubExpRes] -> GBody rep SubExpRes
forall rep res. BodyDec rep -> Stms rep -> [res] -> GBody rep res
Body BodyDec rep
bodydec Stms rep
stms []
  pure $ GPU.Body bodydec stms' res

instance (CSEInOp (op rep)) => CSEInOp (Memory.MemOp op rep) where
  cseInOp :: forall rep. MemOp op rep -> CSEM rep (MemOp op rep)
cseInOp o :: MemOp op rep
o@Memory.Alloc {} = MemOp op rep -> ReaderT (CSEState rep) Identity (MemOp op rep)
forall a. a -> ReaderT (CSEState rep) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MemOp op rep
o
  cseInOp o :: MemOp op rep
o@Memory.EnsureDirect {} = MemOp op rep -> ReaderT (CSEState rep) Identity (MemOp op rep)
forall a. a -> ReaderT (CSEState rep) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MemOp op rep
o
  cseInOp (Memory.Inner op rep
k) = op rep -> MemOp op rep
forall (inner :: * -> *) rep. inner rep -> MemOp inner rep
Memory.Inner (op rep -> MemOp op rep)
-> ReaderT (CSEState rep) Identity (op rep)
-> ReaderT (CSEState rep) Identity (MemOp op rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CSEM (ZonkAny 0) (op rep)
-> ReaderT (CSEState rep) Identity (op rep)
forall rep r otherrep. CSEM rep r -> CSEM otherrep r
subCSE (op rep -> CSEM (ZonkAny 0) (op rep)
forall rep. op rep -> CSEM rep (op rep)
forall op rep. CSEInOp op => op -> CSEM rep op
cseInOp op rep
k)

instance
  (AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
  CSEInOp (SOAC.SOAC (Aliases rep))
  where
  cseInOp :: forall rep. SOAC (Aliases rep) -> CSEM rep (SOAC (Aliases rep))
cseInOp = CSEM (Aliases rep) (SOAC (Aliases rep))
-> CSEM rep (SOAC (Aliases rep))
forall rep r otherrep. CSEM rep r -> CSEM otherrep r
subCSE (CSEM (Aliases rep) (SOAC (Aliases rep))
 -> CSEM rep (SOAC (Aliases rep)))
-> (SOAC (Aliases rep) -> CSEM (Aliases rep) (SOAC (Aliases rep)))
-> SOAC (Aliases rep)
-> CSEM rep (SOAC (Aliases rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOACMapper
  (Aliases rep)
  (Aliases rep)
  (ReaderT (CSEState (Aliases rep)) Identity)
-> SOAC (Aliases rep) -> CSEM (Aliases rep) (SOAC (Aliases rep))
forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
SOAC.mapSOACM ((SubExp -> ReaderT (CSEState (Aliases rep)) Identity SubExp)
-> (Lambda (Aliases rep)
    -> ReaderT
         (CSEState (Aliases rep)) Identity (Lambda (Aliases rep)))
-> (ExtLambda (Aliases rep)
    -> ReaderT
         (CSEState (Aliases rep)) Identity (ExtLambda (Aliases rep)))
-> (VName -> ReaderT (CSEState (Aliases rep)) Identity VName)
-> SOACMapper
     (Aliases rep)
     (Aliases rep)
     (ReaderT (CSEState (Aliases rep)) Identity)
forall frep trep (m :: * -> *).
(SubExp -> m SubExp)
-> (Lambda frep -> m (Lambda trep))
-> (ExtLambda frep -> m (ExtLambda trep))
-> (VName -> m VName)
-> SOACMapper frep trep m
SOAC.SOACMapper SubExp -> ReaderT (CSEState (Aliases rep)) Identity SubExp
forall a. a -> ReaderT (CSEState (Aliases rep)) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda (Aliases rep)
-> ReaderT (CSEState (Aliases rep)) Identity (Lambda (Aliases rep))
forall rep t.
(Aliased rep, CSEInOp (Op rep)) =>
GLambda rep t -> CSEM rep (GLambda rep t)
cseInLambda ExtLambda (Aliases rep)
-> ReaderT
     (CSEState (Aliases rep)) Identity (ExtLambda (Aliases rep))
forall rep t.
(Aliased rep, CSEInOp (Op rep)) =>
GLambda rep t -> CSEM rep (GLambda rep t)
cseInLambda VName -> ReaderT (CSEState (Aliases rep)) Identity VName
forall a. a -> ReaderT (CSEState (Aliases rep)) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure)