{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# 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 qualified Data.Map.Strict as M
import Data.Maybe (isJust)
import Futhark.Analysis.Alias
import Futhark.IR
import Futhark.IR.Aliases
  ( Aliases,
    mkStmsAliases,
    removeFunDefAliases,
    removeProgAliases,
    removeStmAliases,
  )
import qualified Futhark.IR.GPU as GPU
import qualified Futhark.IR.MC as MC
import qualified Futhark.IR.Mem as Memory
import Futhark.IR.Prop.Aliases
import qualified Futhark.IR.SOACS.SOAC as SOAC
import Futhark.Pass
import Futhark.Transform.Substitute

consumedInStms :: Aliased rep => Stms rep -> Names
consumedInStms :: Stms rep -> Names
consumedInStms = ([Names], Names) -> Names
forall a b. (a, b) -> b
snd (([Names], Names) -> Names)
-> (Stms rep -> ([Names], Names)) -> Stms rep -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stms rep -> Result -> ([Names], Names))
-> Result -> Stms rep -> ([Names], Names)
forall a b c. (a -> b -> c) -> b -> a -> c
flip Stms rep -> Result -> ([Names], Names)
forall rep. Aliased rep => Stms rep -> Result -> ([Names], Names)
mkStmsAliases []

-- | 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 ::
  ( ASTRep rep,
    CanBeAliased (Op rep),
    CSEInOp (OpWithAliases (Op rep))
  ) =>
  Bool ->
  Pass rep rep
performCSE :: 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 (Aliases rep) -> Prog rep)
-> PassM (Prog (Aliases rep)) -> PassM (Prog rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Prog (Aliases rep) -> Prog rep
forall rep. CanBeAliased (Op 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 Stms (Aliases rep) -> PassM (Stms (Aliases rep))
forall (f :: * -> *) rep.
(Applicative f, ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
Stms rep -> f (Stms rep)
onConsts Stms (Aliases rep)
-> FunDef (Aliases rep) -> PassM (FunDef (Aliases rep))
forall (f :: * -> *) rep p.
(Applicative f, ASTRep rep, Aliased rep, CSEInOp (Op 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.
(ASTRep rep, CanBeAliased (Op rep)) =>
Prog rep -> Prog (Aliases rep)
aliasAnalysis
  where
    onConsts :: Stms rep -> f (Stms rep)
onConsts Stms rep
stms =
      Stms rep -> f (Stms rep)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms rep -> f (Stms rep)) -> Stms rep -> f (Stms rep)
forall a b. (a -> b) -> a -> b
$
        (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.
(ASTRep rep, 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 (m :: * -> *) a. Monad m => a -> m a
return ()))
            (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 (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.
(ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
Bool -> FunDef rep -> FunDef rep
cseInFunDef Bool
cse_arrays

-- | 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 ::
  ( ASTRep rep,
    CanBeAliased (Op rep),
    CSEInOp (OpWithAliases (Op rep))
  ) =>
  Bool ->
  FunDef rep ->
  FunDef rep
performCSEOnFunDef :: Bool -> FunDef rep -> FunDef rep
performCSEOnFunDef Bool
cse_arrays =
  FunDef (Aliases rep) -> FunDef rep
forall rep.
CanBeAliased (Op 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.
(ASTRep 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.
(ASTRep rep, CanBeAliased (Op rep)) =>
FunDef rep -> FunDef (Aliases rep)
analyseFun

-- | Perform CSE on some statements.
--
-- 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.
performCSEOnStms ::
  ( ASTRep rep,
    CanBeAliased (Op rep),
    CSEInOp (OpWithAliases (Op rep))
  ) =>
  Bool ->
  Stms rep ->
  Stms rep
performCSEOnStms :: Bool -> Stms rep -> Stms rep
performCSEOnStms Bool
cse_arrays =
  (Stm (Aliases rep) -> Stm rep)
-> Seq (Stm (Aliases rep)) -> Stms rep
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Stm (Aliases rep) -> Stm rep
forall rep. CanBeAliased (Op rep) => Stm (Aliases rep) -> Stm rep
removeStmAliases (Seq (Stm (Aliases rep)) -> Stms rep)
-> (Stms rep -> Seq (Stm (Aliases rep))) -> Stms rep -> Stms rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Seq (Stm (Aliases rep)) -> Seq (Stm (Aliases rep))
forall rep.
(ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
Stms rep -> Stms rep
f (Seq (Stm (Aliases rep)) -> Seq (Stm (Aliases rep)))
-> (Stms rep -> Seq (Stm (Aliases rep)))
-> Stms 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)))
-> (Stms rep -> (Seq (Stm (Aliases rep)), AliasesAndConsumed))
-> Stms rep
-> Seq (Stm (Aliases rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AliasTable
-> Stms rep -> (Seq (Stm (Aliases rep)), AliasesAndConsumed)
forall rep.
(ASTRep rep, CanBeAliased (Op 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.
(ASTRep rep, 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 (m :: * -> *) a. Monad m => a -> m a
return ())
          )
          (Bool -> CSEState rep
forall rep. Bool -> CSEState rep
newCSEState Bool
cse_arrays)

cseInFunDef ::
  (ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
  Bool ->
  FunDef rep ->
  FunDef rep
cseInFunDef :: Bool -> FunDef rep -> FunDef rep
cseInFunDef Bool
cse_arrays FunDef rep
fundec =
  FunDef rep
fundec
    { funDefBody :: Body rep
funDefBody =
        Reader (CSEState rep) (Body rep) -> CSEState rep -> Body rep
forall r a. Reader r a -> r -> a
runReader ([Diet] -> Body rep -> Reader (CSEState rep) (Body rep)
forall rep.
(ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
[Diet] -> Body rep -> CSEM rep (Body rep)
cseInBody [Diet]
ds (Body rep -> Reader (CSEState rep) (Body rep))
-> Body rep -> Reader (CSEState rep) (Body rep)
forall a b. (a -> b) -> a -> b
$ FunDef rep -> Body rep
forall rep. FunDef rep -> Body rep
funDefBody FunDef rep
fundec) (CSEState rep -> Body rep) -> CSEState rep -> Body rep
forall a b. (a -> b) -> a -> b
$ Bool -> CSEState rep
forall rep. Bool -> CSEState rep
newCSEState Bool
cse_arrays
    }
  where
    -- XXX: we treat every non-entry result as a consumption here, because we
    -- our core language is not strong enough to fully capture the
    -- aliases we want, so we are turning some parts off (see #803,
    -- #1241, and the related comment in TypeCheck.hs).  This is not a
    -- practical problem while we still perform such aggressive
    -- inlining.
    ds :: [Diet]
ds
      | Maybe EntryPoint -> Bool
forall a. Maybe a -> Bool
isJust (Maybe EntryPoint -> Bool) -> Maybe EntryPoint -> Bool
forall a b. (a -> b) -> a -> b
$ FunDef rep -> Maybe EntryPoint
forall rep. FunDef rep -> Maybe EntryPoint
funDefEntryPoint FunDef rep
fundec = (RetType rep -> Diet) -> [RetType rep] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase ExtShape Uniqueness -> Diet
forall shape. TypeBase shape Uniqueness -> Diet
diet (TypeBase ExtShape Uniqueness -> Diet)
-> (RetType rep -> TypeBase ExtShape Uniqueness)
-> RetType rep
-> Diet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RetType rep -> TypeBase ExtShape Uniqueness
forall t. DeclExtTyped t => t -> TypeBase ExtShape Uniqueness
declExtTypeOf) ([RetType rep] -> [Diet]) -> [RetType rep] -> [Diet]
forall a b. (a -> b) -> a -> b
$ FunDef rep -> [RetType rep]
forall rep. FunDef rep -> [RetType rep]
funDefRetType FunDef rep
fundec
      | Bool
otherwise = (RetType rep -> Diet) -> [RetType rep] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map RetType rep -> Diet
forall t. DeclExtTyped t => t -> Diet
retDiet ([RetType rep] -> [Diet]) -> [RetType rep] -> [Diet]
forall a b. (a -> b) -> a -> b
$ FunDef rep -> [RetType rep]
forall rep. FunDef rep -> [RetType rep]
funDefRetType FunDef rep
fundec
    retDiet :: t -> Diet
retDiet t
t
      | TypeBase ExtShape Uniqueness -> Bool
forall shape u. TypeBase shape u -> Bool
primType (TypeBase ExtShape Uniqueness -> Bool)
-> TypeBase ExtShape Uniqueness -> Bool
forall a b. (a -> b) -> a -> b
$ t -> TypeBase ExtShape Uniqueness
forall t. DeclExtTyped t => t -> TypeBase ExtShape Uniqueness
declExtTypeOf t
t = Diet
Observe
      | Bool
otherwise = Diet
Consume

type CSEM rep = Reader (CSEState rep)

cseInBody ::
  (ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
  [Diet] ->
  Body rep ->
  CSEM rep (Body rep)
cseInBody :: [Diet] -> Body rep -> CSEM rep (Body rep)
cseInBody [Diet]
ds (Body BodyDec rep
bodydec Stms rep
stms Result
res) = do
  (Stms rep
stms', Result
res') <-
    Names
-> [Stm rep] -> CSEM rep Result -> CSEM rep (Stms rep, Result)
forall rep a.
(ASTRep rep, 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 Result -> CSEM rep (Stms rep, Result))
-> CSEM rep Result -> CSEM rep (Stms rep, Result)
forall a b. (a -> b) -> a -> b
$ do
      CSEState (ExpressionSubstitutions rep
_, NameSubstitutions
nsubsts) Bool
_ <- ReaderT (CSEState rep) Identity (CSEState rep)
forall r (m :: * -> *). MonadReader r m => m r
ask
      Result -> CSEM rep Result
forall (m :: * -> *) a. Monad m => a -> m a
return (Result -> CSEM rep Result) -> Result -> CSEM rep Result
forall a b. (a -> b) -> a -> b
$ NameSubstitutions -> Result -> Result
forall a. Substitute a => NameSubstitutions -> a -> a
substituteNames NameSubstitutions
nsubsts Result
res
  Body rep -> CSEM rep (Body rep)
forall (m :: * -> *) a. Monad m => a -> m a
return (Body rep -> CSEM rep (Body rep))
-> Body rep -> CSEM rep (Body rep)
forall a b. (a -> b) -> a -> b
$ BodyDec rep -> Stms rep -> Result -> Body rep
forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body BodyDec rep
bodydec Stms rep
stms' Result
res'
  where
    ([Names]
res_als, Names
stms_cons) = Stms rep -> Result -> ([Names], Names)
forall rep. Aliased rep => Stms rep -> Result -> ([Names], Names)
mkStmsAliases Stms rep
stms Result
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 ::
  (ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
  Lambda rep ->
  CSEM rep (Lambda rep)
cseInLambda :: Lambda rep -> CSEM rep (Lambda rep)
cseInLambda Lambda rep
lam = do
  Body rep
body' <- [Diet] -> Body rep -> CSEM rep (Body rep)
forall rep.
(ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
[Diet] -> Body rep -> CSEM rep (Body rep)
cseInBody ((Type -> Diet) -> [Type] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map (Diet -> Type -> Diet
forall a b. a -> b -> a
const Diet
Observe) ([Type] -> [Diet]) -> [Type] -> [Diet]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Type]
forall rep. Lambda rep -> [Type]
lambdaReturnType Lambda rep
lam) (Body rep -> CSEM rep (Body rep))
-> Body rep -> CSEM rep (Body rep)
forall a b. (a -> b) -> a -> b
$ Lambda rep -> Body rep
forall rep. Lambda rep -> Body rep
lambdaBody Lambda rep
lam
  Lambda rep -> CSEM rep (Lambda rep)
forall (m :: * -> *) a. Monad m => a -> m a
return Lambda rep
lam {lambdaBody :: Body rep
lambdaBody = Body rep
body'}

cseInStms ::
  (ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
  Names ->
  [Stm rep] ->
  CSEM rep a ->
  CSEM rep (Stms rep, a)
cseInStms :: Names -> [Stm rep] -> CSEM rep a -> CSEM rep (Stms rep, a)
cseInStms Names
_ [] CSEM rep a
m = do
  a
a <- CSEM rep a
m
  (Stms rep, a) -> CSEM rep (Stms rep, a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Stms rep
forall a. Monoid a => a
mempty, a
a)
cseInStms Names
consumed (Stm rep
stm : [Stm rep]
stms) CSEM rep a
m =
  Names
-> Stm rep
-> ([Stm rep] -> CSEM rep (Stms rep, a))
-> CSEM rep (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] -> CSEM rep (Stms rep, a)) -> CSEM rep (Stms rep, a))
-> ([Stm rep] -> CSEM rep (Stms rep, a)) -> CSEM rep (Stms rep, a)
forall a b. (a -> b) -> a -> b
$ \[Stm rep]
stm' -> do
    (Stms rep
stms', a
a) <- Names -> [Stm rep] -> CSEM rep a -> CSEM rep (Stms rep, a)
forall rep a.
(ASTRep rep, 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 rep]
stm'' <- (Stm rep -> ReaderT (CSEState rep) Identity (Stm rep))
-> [Stm rep] -> ReaderT (CSEState rep) Identity [Stm rep]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Stm rep -> ReaderT (CSEState rep) Identity (Stm rep)
forall rep.
(ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
Stm rep -> ReaderT (CSEState rep) Identity (Stm rep)
nestedCSE [Stm rep]
stm'
    (Stms rep, a) -> CSEM rep (Stms rep, a)
forall (m :: * -> *) a. Monad m => a -> m a
return ([Stm rep] -> Stms rep
forall rep. [Stm rep] -> Stms rep
stmsFromList [Stm rep]
stm'' Stms rep -> Stms rep -> Stms rep
forall a. Semigroup a => a -> a -> a
<> Stms rep
stms', a
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
              DoLoop [(FParam rep, SubExp)]
merge LoopForm rep
_ Body rep
_ -> ((FParam rep, SubExp) -> Diet) -> [(FParam rep, SubExp)] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase Shape Uniqueness -> Diet
forall shape. TypeBase shape Uniqueness -> Diet
diet (TypeBase Shape Uniqueness -> Diet)
-> ((FParam rep, SubExp) -> TypeBase Shape Uniqueness)
-> (FParam rep, SubExp)
-> Diet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FParam rep -> TypeBase Shape Uniqueness
forall t. DeclTyped t => t -> TypeBase Shape Uniqueness
declTypeOf (FParam rep -> TypeBase Shape Uniqueness)
-> ((FParam rep, SubExp) -> FParam rep)
-> (FParam rep, SubExp)
-> TypeBase Shape Uniqueness
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FParam rep, SubExp) -> FParam rep
forall a b. (a, b) -> a
fst) [(FParam 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'
      Exp rep
e <- Mapper rep rep (ReaderT (CSEState rep) Identity)
-> Exp rep -> ReaderT (CSEState rep) Identity (Exp rep)
forall (m :: * -> *) frep trep.
(Applicative m, Monad m) =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM ([Diet] -> Mapper rep rep (ReaderT (CSEState rep) Identity)
forall trep.
(ASTRep trep, Aliased trep, CSEInOp (Op trep)) =>
[Diet] -> Mapper trep trep (ReaderT (CSEState trep) 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'
      Stm rep -> ReaderT (CSEState rep) Identity (Stm rep)
forall (m :: * -> *) a. Monad m => a -> m a
return Stm rep
stm' {stmExp :: Exp rep
stmExp = Exp rep
e}

    cse :: [Diet] -> Mapper trep trep (ReaderT (CSEState trep) Identity)
cse [Diet]
ds =
      Mapper trep trep (ReaderT (CSEState trep) Identity)
forall (m :: * -> *) rep. Monad m => Mapper rep rep m
identityMapper
        { mapOnBody :: Scope trep
-> Body trep -> ReaderT (CSEState trep) Identity (Body trep)
mapOnBody = (Body trep -> ReaderT (CSEState trep) Identity (Body trep))
-> Scope trep
-> Body trep
-> ReaderT (CSEState trep) Identity (Body trep)
forall a b. a -> b -> a
const ((Body trep -> ReaderT (CSEState trep) Identity (Body trep))
 -> Scope trep
 -> Body trep
 -> ReaderT (CSEState trep) Identity (Body trep))
-> (Body trep -> ReaderT (CSEState trep) Identity (Body trep))
-> Scope trep
-> Body trep
-> ReaderT (CSEState trep) Identity (Body trep)
forall a b. (a -> b) -> a -> b
$ [Diet] -> Body trep -> ReaderT (CSEState trep) Identity (Body trep)
forall rep.
(ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
[Diet] -> Body rep -> CSEM rep (Body rep)
cseInBody [Diet]
ds,
          mapOnOp :: Op trep -> ReaderT (CSEState trep) Identity (Op trep)
mapOnOp = Op trep -> ReaderT (CSEState trep) Identity (Op trep)
forall op rep. CSEInOp op => op -> CSEM rep op
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 :: Names -> Stm rep -> ([Stm rep] -> CSEM rep a) -> CSEM rep a
cseInStm Names
consumed (Let Pat (LetDec rep)
pat (StmAux Certs
cs Attrs
attrs ExpDec rep
edec) Exp rep
e) [Stm rep] -> CSEM rep a
m = do
  CSEState (ExpressionSubstitutions rep
esubsts, NameSubstitutions
nsubsts) Bool
cse_arrays <- ReaderT (CSEState rep) Identity (CSEState rep)
forall r (m :: * -> *). MonadReader r m => m r
ask
  let e' :: Exp rep
e' = NameSubstitutions -> Exp rep -> Exp rep
forall a. Substitute a => NameSubstitutions -> a -> a
substituteNames NameSubstitutions
nsubsts Exp rep
e
      pat' :: Pat (LetDec rep)
pat' = NameSubstitutions -> Pat (LetDec rep) -> Pat (LetDec rep)
forall a. Substitute a => NameSubstitutions -> a -> a
substituteNames NameSubstitutions
nsubsts Pat (LetDec rep)
pat
  if (PatElem (LetDec rep) -> Bool) -> [PatElem (LetDec rep)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Bool -> PatElem (LetDec rep) -> Bool
forall dec. Typed dec => Bool -> PatElem dec -> Bool
bad Bool
cse_arrays) ([PatElem (LetDec rep)] -> Bool) -> [PatElem (LetDec rep)] -> Bool
forall a b. (a -> b) -> a -> b
$ Pat (LetDec rep) -> [PatElem (LetDec rep)]
forall dec. Pat dec -> [PatElem dec]
patElems Pat (LetDec rep)
pat
    then [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' (Certs -> Attrs -> ExpDec rep -> StmAux (ExpDec rep)
forall dec. Certs -> Attrs -> dec -> StmAux dec
StmAux Certs
cs Attrs
attrs ExpDec rep
edec) Exp rep
e']
    else case (ExpDec rep, Exp rep)
-> ExpressionSubstitutions rep -> Maybe (Pat (LetDec rep))
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (ExpDec rep
edec, Exp rep
e') ExpressionSubstitutions rep
esubsts of
      Just Pat (LetDec rep)
subpat ->
        (CSEState rep -> CSEState rep) -> CSEM rep a -> CSEM rep a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (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) (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']) (Certs -> Attrs -> ExpDec rep -> StmAux (ExpDec rep)
forall dec. Certs -> Attrs -> dec -> StmAux dec
StmAux Certs
cs Attrs
attrs ExpDec rep
edec) (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$
                    BasicOp -> Exp rep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ 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 :: VName
patElemName = VName
name}
                ]
          [Stm rep] -> CSEM rep a
m [Stm rep]
lets
      Maybe (Pat (LetDec rep))
_ ->
        (CSEState rep -> CSEState rep) -> CSEM rep a -> CSEM rep a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (Pat (LetDec rep)
-> ExpDec rep -> Exp rep -> CSEState rep -> CSEState rep
forall rep.
ASTRep rep =>
Pat (LetDec rep)
-> ExpDec rep -> Exp rep -> CSEState rep -> CSEState rep
addExpSubst Pat (LetDec rep)
pat' ExpDec rep
edec 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' (Certs -> Attrs -> ExpDec rep -> StmAux (ExpDec rep)
forall dec. Certs -> Attrs -> dec -> StmAux dec
StmAux Certs
cs Attrs
attrs ExpDec rep
edec) Exp rep
e']
  where
    bad :: Bool -> PatElem dec -> Bool
bad Bool
cse_arrays PatElem dec
pe
      | Mem {} <- PatElem dec -> Type
forall dec. Typed dec => PatElem dec -> Type
patElemType PatElem dec
pe = Bool
True
      | Array {} <- PatElem dec -> Type
forall dec. Typed dec => PatElem dec -> Type
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)
    (Pat (LetDec rep))

type NameSubstitutions = M.Map VName VName

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

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

mkSubsts :: Pat dec -> Pat dec -> M.Map VName VName
mkSubsts :: Pat dec -> Pat dec -> NameSubstitutions
mkSubsts Pat dec
pat Pat dec
vs = [(VName, VName)] -> NameSubstitutions
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, VName)] -> NameSubstitutions)
-> [(VName, VName)] -> NameSubstitutions
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 :: Pat dec -> Pat dec -> CSEState rep -> CSEState rep
addNameSubst Pat dec
pat Pat dec
subpat (CSEState (ExpressionSubstitutions rep
esubsts, NameSubstitutions
nsubsts) Bool
cse_arrays) =
  (ExpressionSubstitutions rep, NameSubstitutions)
-> Bool -> CSEState rep
forall rep.
(ExpressionSubstitutions rep, NameSubstitutions)
-> Bool -> CSEState rep
CSEState (ExpressionSubstitutions rep
esubsts, Pat dec -> Pat dec -> NameSubstitutions
forall dec. Pat dec -> Pat dec -> NameSubstitutions
mkSubsts Pat dec
pat Pat dec
subpat NameSubstitutions -> NameSubstitutions -> NameSubstitutions
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` NameSubstitutions
nsubsts) Bool
cse_arrays

addExpSubst ::
  ASTRep rep =>
  Pat (LetDec rep) ->
  ExpDec rep ->
  Exp rep ->
  CSEState rep ->
  CSEState rep
addExpSubst :: Pat (LetDec rep)
-> ExpDec rep -> Exp rep -> CSEState rep -> CSEState rep
addExpSubst Pat (LetDec rep)
pat ExpDec rep
edec Exp rep
e (CSEState (ExpressionSubstitutions rep
esubsts, NameSubstitutions
nsubsts) Bool
cse_arrays) =
  (ExpressionSubstitutions rep, NameSubstitutions)
-> Bool -> CSEState rep
forall rep.
(ExpressionSubstitutions rep, NameSubstitutions)
-> Bool -> CSEState rep
CSEState ((ExpDec rep, Exp rep)
-> 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) Pat (LetDec rep)
pat ExpressionSubstitutions rep
esubsts, NameSubstitutions
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 () where
  cseInOp :: () -> CSEM rep ()
cseInOp () = () -> CSEM rep ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

subCSE :: CSEM rep r -> CSEM otherrep r
subCSE :: CSEM rep r -> CSEM otherrep r
subCSE CSEM rep r
m = do
  CSEState (ExpressionSubstitutions otherrep, NameSubstitutions)
_ Bool
cse_arrays <- ReaderT (CSEState otherrep) Identity (CSEState otherrep)
forall r (m :: * -> *). MonadReader r m => m r
ask
  r -> CSEM otherrep r
forall (m :: * -> *) a. Monad m => a -> m a
return (r -> CSEM otherrep r) -> r -> CSEM otherrep r
forall a b. (a -> b) -> a -> b
$ CSEM rep r -> CSEState rep -> r
forall r a. Reader r a -> r -> a
runReader CSEM rep r
m (CSEState rep -> r) -> CSEState rep -> r
forall a b. (a -> b) -> a -> b
$ Bool -> CSEState rep
forall rep. Bool -> CSEState rep
newCSEState Bool
cse_arrays

instance
  ( ASTRep rep,
    Aliased rep,
    CSEInOp (Op rep),
    CSEInOp op
  ) =>
  CSEInOp (GPU.HostOp rep op)
  where
  cseInOp :: HostOp rep op -> CSEM rep (HostOp rep op)
cseInOp (GPU.SegOp SegOp SegLevel rep
op) = SegOp SegLevel rep -> HostOp rep op
forall rep op. SegOp SegLevel rep -> HostOp rep op
GPU.SegOp (SegOp SegLevel rep -> HostOp rep op)
-> ReaderT (CSEState rep) Identity (SegOp SegLevel rep)
-> CSEM rep (HostOp rep op)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegOp SegLevel rep
-> ReaderT (CSEState rep) Identity (SegOp SegLevel rep)
forall op rep. CSEInOp op => op -> CSEM rep op
cseInOp SegOp SegLevel rep
op
  cseInOp (GPU.OtherOp op
op) = op -> HostOp rep op
forall rep op. op -> HostOp rep op
GPU.OtherOp (op -> HostOp rep op)
-> ReaderT (CSEState rep) Identity op -> CSEM rep (HostOp rep op)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> op -> ReaderT (CSEState rep) Identity op
forall op rep. CSEInOp op => op -> CSEM rep op
cseInOp op
op
  cseInOp HostOp rep op
x = HostOp rep op -> CSEM rep (HostOp rep op)
forall (m :: * -> *) a. Monad m => a -> m a
return HostOp rep op
x

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

instance
  (ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
  CSEInOp (GPU.SegOp lvl rep)
  where
  cseInOp :: 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.
(Applicative m, 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))
-> (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))
-> (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 (m :: * -> *) a. Monad m => a -> m a
return Lambda rep -> ReaderT (CSEState rep) Identity (Lambda rep)
forall rep.
(ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
Lambda rep -> CSEM rep (Lambda rep)
cseInLambda KernelBody rep -> ReaderT (CSEState rep) Identity (KernelBody rep)
forall rep.
(ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
KernelBody rep -> CSEM rep (KernelBody rep)
cseInKernelBody VName -> ReaderT (CSEState rep) Identity VName
forall (m :: * -> *) a. Monad m => a -> m a
return lvl -> ReaderT (CSEState rep) Identity lvl
forall (m :: * -> *) a. Monad m => a -> m a
return)

cseInKernelBody ::
  (ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
  GPU.KernelBody rep ->
  CSEM rep (GPU.KernelBody rep)
cseInKernelBody :: KernelBody rep -> CSEM rep (KernelBody rep)
cseInKernelBody (GPU.KernelBody BodyDec rep
bodydec Stms rep
stms [KernelResult]
res) = do
  Body BodyDec rep
_ Stms rep
stms' Result
_ <- [Diet] -> Body rep -> CSEM rep (Body rep)
forall rep.
(ASTRep 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) (Body rep -> CSEM rep (Body rep))
-> Body rep -> CSEM rep (Body rep)
forall a b. (a -> b) -> a -> b
$ BodyDec rep -> Stms rep -> Result -> Body rep
forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body BodyDec rep
bodydec Stms rep
stms []
  KernelBody rep -> CSEM rep (KernelBody rep)
forall (m :: * -> *) a. Monad m => a -> m a
return (KernelBody rep -> CSEM rep (KernelBody rep))
-> KernelBody rep -> CSEM rep (KernelBody rep)
forall a b. (a -> b) -> a -> b
$ BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
forall rep.
BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
GPU.KernelBody BodyDec rep
bodydec Stms rep
stms' [KernelResult]
res

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

instance
  ( ASTRep rep,
    CanBeAliased (Op rep),
    CSEInOp (OpWithAliases (Op rep))
  ) =>
  CSEInOp (SOAC.SOAC (Aliases rep))
  where
  cseInOp :: 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.
(Applicative m, 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)))
-> (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))
-> (VName -> m VName)
-> SOACMapper frep trep m
SOAC.SOACMapper SubExp -> ReaderT (CSEState (Aliases rep)) Identity SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return Lambda (Aliases rep)
-> ReaderT (CSEState (Aliases rep)) Identity (Lambda (Aliases rep))
forall rep.
(ASTRep rep, Aliased rep, CSEInOp (Op rep)) =>
Lambda rep -> CSEM rep (Lambda rep)
cseInLambda VName -> ReaderT (CSEState (Aliases rep)) Identity VName
forall (m :: * -> *) a. Monad m => a -> m a
return)