module Futhark.IR.Traversals
(
Mapper (..),
identityMapper,
mapExpM,
mapExp,
Walker (..),
identityWalker,
walkExpM,
TraverseOpStms (..),
OpStmsTraverser,
traverseLambdaStms,
)
where
import Control.Monad
import Control.Monad.Identity
import Data.Bitraversable
import Data.Foldable (traverse_)
import Data.List.NonEmpty (NonEmpty (..))
import Futhark.IR.Prop.Scope
import Futhark.IR.Prop.Types (mapOnType)
import Futhark.IR.Syntax
data Mapper frep trep m = Mapper
{ forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp :: SubExp -> m SubExp,
forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> Scope trep -> Body frep -> m (Body trep)
mapOnBody :: Scope trep -> Body frep -> m (Body trep),
forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName :: VName -> m VName,
forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> RetType frep -> m (RetType trep)
mapOnRetType :: RetType frep -> m (RetType trep),
forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> BranchType frep -> m (BranchType trep)
mapOnBranchType :: BranchType frep -> m (BranchType trep),
forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> FParam frep -> m (FParam trep)
mapOnFParam :: FParam frep -> m (FParam trep),
forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> LParam frep -> m (LParam trep)
mapOnLParam :: LParam frep -> m (LParam trep),
forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> Op frep -> m (Op trep)
mapOnOp :: Op frep -> m (Op trep)
}
identityMapper :: Monad m => Mapper rep rep m
identityMapper :: forall {k} (m :: * -> *) (rep :: k). Monad m => Mapper rep rep m
identityMapper =
Mapper
{ mapOnSubExp :: SubExp -> m SubExp
mapOnSubExp = forall (f :: * -> *) a. Applicative f => a -> f a
pure,
mapOnBody :: Scope rep -> Body rep -> m (Body rep)
mapOnBody = forall a b. a -> b -> a
const forall (f :: * -> *) a. Applicative f => a -> f a
pure,
mapOnVName :: VName -> m VName
mapOnVName = forall (f :: * -> *) a. Applicative f => a -> f a
pure,
mapOnRetType :: RetType rep -> m (RetType rep)
mapOnRetType = forall (f :: * -> *) a. Applicative f => a -> f a
pure,
mapOnBranchType :: BranchType rep -> m (BranchType rep)
mapOnBranchType = forall (f :: * -> *) a. Applicative f => a -> f a
pure,
mapOnFParam :: FParam rep -> m (FParam rep)
mapOnFParam = forall (f :: * -> *) a. Applicative f => a -> f a
pure,
mapOnLParam :: LParam rep -> m (LParam rep)
mapOnLParam = forall (f :: * -> *) a. Applicative f => a -> f a
pure,
mapOnOp :: Op rep -> m (Op rep)
mapOnOp = forall (f :: * -> *) a. Applicative f => a -> f a
pure
}
mapExpM ::
Monad m =>
Mapper frep trep m ->
Exp frep ->
m (Exp trep)
mapExpM :: forall {k} {k} (m :: * -> *) (frep :: k) (trep :: k).
Monad m =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper frep trep m
tv (BasicOp (SubExp SubExp
se)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> BasicOp
SubExp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
se)
mapExpM Mapper frep trep m
tv (BasicOp (ArrayLit [SubExp]
els Type
rowt)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( [SubExp] -> Type -> BasicOp
ArrayLit
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) [SubExp]
els
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (m :: * -> *) u.
Monad m =>
(SubExp -> m SubExp) -> TypeBase Shape u -> m (TypeBase Shape u)
mapOnType (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) Type
rowt
)
mapExpM Mapper frep trep m
tv (BasicOp (BinOp BinOp
bop SubExp
x SubExp
y)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (BinOp -> SubExp -> SubExp -> BasicOp
BinOp BinOp
bop forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
y)
mapExpM Mapper frep trep m
tv (BasicOp (CmpOp CmpOp
op SubExp
x SubExp
y)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp CmpOp
op forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
y)
mapExpM Mapper frep trep m
tv (BasicOp (ConvOp ConvOp
conv SubExp
x)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ConvOp -> SubExp -> BasicOp
ConvOp ConvOp
conv forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
x)
mapExpM Mapper frep trep m
tv (BasicOp (UnOp UnOp
op SubExp
x)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (UnOp -> SubExp -> BasicOp
UnOp UnOp
op forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
x)
mapExpM Mapper frep trep m
tv (Match [SubExp]
ses [Case (Body frep)]
cases Body frep
defbody (MatchDec [BranchType frep]
ts MatchSort
s)) =
forall {k} (rep :: k).
[SubExp]
-> [Case (Body rep)]
-> Body rep
-> MatchDec (BranchType rep)
-> Exp rep
Match
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) [SubExp]
ses
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Case (Body frep) -> m (Case (Body trep))
mapOnCase [Case (Body frep)]
cases
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> Scope trep -> Body frep -> m (Body trep)
mapOnBody Mapper frep trep m
tv forall a. Monoid a => a
mempty Body frep
defbody
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (forall rt. [rt] -> MatchSort -> MatchDec rt
MatchDec forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> BranchType frep -> m (BranchType trep)
mapOnBranchType Mapper frep trep m
tv) [BranchType frep]
ts forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure MatchSort
s)
where
mapOnCase :: Case (Body frep) -> m (Case (Body trep))
mapOnCase (Case [Maybe PrimValue]
vs Body frep
body) = forall body. [Maybe PrimValue] -> body -> Case body
Case [Maybe PrimValue]
vs forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> Scope trep -> Body frep -> m (Body trep)
mapOnBody Mapper frep trep m
tv forall a. Monoid a => a
mempty Body frep
body
mapExpM Mapper frep trep m
tv (Apply Name
fname [(SubExp, Diet)]
args [RetType frep]
ret (Safety, SrcLoc, [SrcLoc])
loc) = do
[(SubExp, Diet)]
args' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(SubExp, Diet)]
args forall a b. (a -> b) -> a -> b
$ \(SubExp
arg, Diet
d) ->
(,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
arg forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure Diet
d
forall {k} (rep :: k).
Name
-> [(SubExp, Diet)]
-> [RetType rep]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp rep
Apply Name
fname [(SubExp, Diet)]
args' forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> RetType frep -> m (RetType trep)
mapOnRetType Mapper frep trep m
tv) [RetType frep]
ret forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Safety, SrcLoc, [SrcLoc])
loc
mapExpM Mapper frep trep m
tv (BasicOp (Index VName
arr Slice SubExp
slice)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( VName -> Slice SubExp -> BasicOp
Index
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
arr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) Slice SubExp
slice
)
mapExpM Mapper frep trep m
tv (BasicOp (Update Safety
safety VName
arr Slice SubExp
slice SubExp
se)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( Safety -> VName -> Slice SubExp -> SubExp -> BasicOp
Update Safety
safety
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
arr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) Slice SubExp
slice
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
se
)
mapExpM Mapper frep trep m
tv (BasicOp (FlatIndex VName
arr FlatSlice SubExp
slice)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( VName -> FlatSlice SubExp -> BasicOp
FlatIndex
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
arr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) FlatSlice SubExp
slice
)
mapExpM Mapper frep trep m
tv (BasicOp (FlatUpdate VName
arr FlatSlice SubExp
slice VName
se)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( VName -> FlatSlice SubExp -> VName -> BasicOp
FlatUpdate
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
arr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) FlatSlice SubExp
slice
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
se
)
mapExpM Mapper frep trep m
tv (BasicOp (Iota SubExp
n SubExp
x SubExp
s IntType
et)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> SubExp -> SubExp -> IntType -> BasicOp
Iota forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
n forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
s forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure IntType
et)
mapExpM Mapper frep trep m
tv (BasicOp (Replicate Shape
shape SubExp
vexp)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Shape -> SubExp -> BasicOp
Replicate forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (m :: * -> *) (frep :: k) (trep :: k).
Monad m =>
Mapper frep trep m -> Shape -> m Shape
mapOnShape Mapper frep trep m
tv Shape
shape forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
vexp)
mapExpM Mapper frep trep m
tv (BasicOp (Scratch PrimType
t [SubExp]
shape)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PrimType -> [SubExp] -> BasicOp
Scratch PrimType
t forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) [SubExp]
shape)
mapExpM Mapper frep trep m
tv (BasicOp (Reshape ReshapeKind
kind Shape
shape VName
arrexp)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( ReshapeKind -> Shape -> VName -> BasicOp
Reshape ReshapeKind
kind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) Shape
shape
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
arrexp
)
mapExpM Mapper frep trep m
tv (BasicOp (Rearrange [Int]
perm VName
e)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Int] -> VName -> BasicOp
Rearrange [Int]
perm forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
e)
mapExpM Mapper frep trep m
tv (BasicOp (Rotate [SubExp]
es VName
e)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([SubExp] -> VName -> BasicOp
Rotate forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) [SubExp]
es forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
e)
mapExpM Mapper frep trep m
tv (BasicOp (Concat Int
i (VName
x :| [VName]
ys) SubExp
size)) = do
VName
x' <- forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
x
[VName]
ys' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv) [VName]
ys
SubExp
size' <- forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
size
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ Int -> NonEmpty VName -> SubExp -> BasicOp
Concat Int
i (VName
x' forall a. a -> [a] -> NonEmpty a
:| [VName]
ys') SubExp
size'
mapExpM Mapper frep trep m
tv (BasicOp (Copy VName
e)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> BasicOp
Copy forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
e)
mapExpM Mapper frep trep m
tv (BasicOp (Manifest [Int]
perm VName
e)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Int] -> VName -> BasicOp
Manifest [Int]
perm forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
e)
mapExpM Mapper frep trep m
tv (BasicOp (Assert SubExp
e ErrorMsg SubExp
msg (SrcLoc, [SrcLoc])
loc)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> ErrorMsg SubExp -> (SrcLoc, [SrcLoc]) -> BasicOp
Assert forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
e forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) ErrorMsg SubExp
msg forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure (SrcLoc, [SrcLoc])
loc)
mapExpM Mapper frep trep m
tv (BasicOp (Opaque OpaqueOp
op SubExp
e)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (OpaqueOp -> SubExp -> BasicOp
Opaque OpaqueOp
op forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
e)
mapExpM Mapper frep trep m
tv (BasicOp (UpdateAcc VName
v [SubExp]
is [SubExp]
ses)) =
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( VName -> [SubExp] -> [SubExp] -> BasicOp
UpdateAcc
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
v
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) [SubExp]
is
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) [SubExp]
ses
)
mapExpM Mapper frep trep m
tv (WithAcc [WithAccInput frep]
inputs Lambda frep
lam) =
forall {k} (rep :: k). [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {t :: * -> * -> *} {t :: * -> *} {t :: * -> *}
{t :: * -> *}.
(Bitraversable t, Traversable t, Traversable t, Traversable t) =>
(Shape, t VName, t (t (Lambda frep) (t SubExp)))
-> m (Shape, t VName, t (t (Lambda trep) (t SubExp)))
onInput [WithAccInput frep]
inputs forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {k} {k} (m :: * -> *) (frep :: k) (trep :: k).
Monad m =>
Mapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnLambda Mapper frep trep m
tv Lambda frep
lam
where
onInput :: (Shape, t VName, t (t (Lambda frep) (t SubExp)))
-> m (Shape, t VName, t (t (Lambda trep) (t SubExp)))
onInput (Shape
shape, t VName
vs, t (t (Lambda frep) (t SubExp))
op) =
(,,)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (m :: * -> *) (frep :: k) (trep :: k).
Monad m =>
Mapper frep trep m -> Shape -> m Shape
mapOnShape Mapper frep trep m
tv Shape
shape
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv) t VName
vs
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse (forall {k} {k} (m :: * -> *) (frep :: k) (trep :: k).
Monad m =>
Mapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnLambda Mapper frep trep m
tv) (forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv))) t (t (Lambda frep) (t SubExp))
op
mapExpM Mapper frep trep m
tv (DoLoop [(FParam frep, SubExp)]
merge LoopForm frep
form Body frep
loopbody) = do
[Param (FParamInfo trep)]
params' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> FParam frep -> m (FParam trep)
mapOnFParam Mapper frep trep m
tv) [FParam frep]
params
LoopForm trep
form' <- forall {k} {k} (m :: * -> *) (frep :: k) (trep :: k).
Monad m =>
Mapper frep trep m -> LoopForm frep -> m (LoopForm trep)
mapOnLoopForm Mapper frep trep m
tv LoopForm frep
form
let scope :: Scope trep
scope = forall {k} (rep :: k) a. Scoped rep a => a -> Scope rep
scopeOf LoopForm trep
form' forall a. Semigroup a => a -> a -> a
<> forall {k} (rep :: k) dec.
(FParamInfo rep ~ dec) =>
[Param dec] -> Scope rep
scopeOfFParams [Param (FParamInfo trep)]
params'
forall {k} (rep :: k).
[(FParam rep, SubExp)] -> LoopForm rep -> Body rep -> Exp rep
DoLoop
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall a b. [a] -> [b] -> [(a, b)]
zip [Param (FParamInfo trep)]
params' forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) [SubExp]
args)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure LoopForm trep
form'
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> Scope trep -> Body frep -> m (Body trep)
mapOnBody Mapper frep trep m
tv Scope trep
scope Body frep
loopbody
where
([FParam frep]
params, [SubExp]
args) = forall a b. [(a, b)] -> ([a], [b])
unzip [(FParam frep, SubExp)]
merge
mapExpM Mapper frep trep m
tv (Op Op frep
op) =
forall {k} (rep :: k). Op rep -> Exp rep
Op forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> Op frep -> m (Op trep)
mapOnOp Mapper frep trep m
tv Op frep
op
mapOnShape :: Monad m => Mapper frep trep m -> Shape -> m Shape
mapOnShape :: forall {k} {k} (m :: * -> *) (frep :: k) (trep :: k).
Monad m =>
Mapper frep trep m -> Shape -> m Shape
mapOnShape Mapper frep trep m
tv (Shape [SubExp]
ds) = forall d. [d] -> ShapeBase d
Shape forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv) [SubExp]
ds
mapOnLoopForm ::
Monad m =>
Mapper frep trep m ->
LoopForm frep ->
m (LoopForm trep)
mapOnLoopForm :: forall {k} {k} (m :: * -> *) (frep :: k) (trep :: k).
Monad m =>
Mapper frep trep m -> LoopForm frep -> m (LoopForm trep)
mapOnLoopForm Mapper frep trep m
tv (ForLoop VName
i IntType
it SubExp
bound [(LParam frep, VName)]
loop_vars) =
forall {k} (rep :: k).
VName -> IntType -> SubExp -> [(LParam rep, VName)] -> LoopForm rep
ForLoop
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
i
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure IntType
it
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv SubExp
bound
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (forall a b. [a] -> [b] -> [(a, b)]
zip forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> LParam frep -> m (LParam trep)
mapOnLParam Mapper frep trep m
tv) [LParam frep]
loop_lparams forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv) [VName]
loop_arrs)
where
([LParam frep]
loop_lparams, [VName]
loop_arrs) = forall a b. [(a, b)] -> ([a], [b])
unzip [(LParam frep, VName)]
loop_vars
mapOnLoopForm Mapper frep trep m
tv (WhileLoop VName
cond) =
forall {k} (rep :: k). VName -> LoopForm rep
WhileLoop forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> VName -> m VName
mapOnVName Mapper frep trep m
tv VName
cond
mapOnLambda ::
Monad m =>
Mapper frep trep m ->
Lambda frep ->
m (Lambda trep)
mapOnLambda :: forall {k} {k} (m :: * -> *) (frep :: k) (trep :: k).
Monad m =>
Mapper frep trep m -> Lambda frep -> m (Lambda trep)
mapOnLambda Mapper frep trep m
tv (Lambda [LParam frep]
params Body frep
body [Type]
ret) = do
[Param (LParamInfo trep)]
params' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> LParam frep -> m (LParam trep)
mapOnLParam Mapper frep trep m
tv) [LParam frep]
params
forall {k} (rep :: k).
[LParam rep] -> Body rep -> [Type] -> Lambda rep
Lambda [Param (LParamInfo trep)]
params'
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> Scope trep -> Body frep -> m (Body trep)
mapOnBody Mapper frep trep m
tv (forall {k} (rep :: k) dec.
(LParamInfo rep ~ dec) =>
[Param dec] -> Scope rep
scopeOfLParams [Param (LParamInfo trep)]
params') Body frep
body
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *) u.
Monad m =>
(SubExp -> m SubExp) -> TypeBase Shape u -> m (TypeBase Shape u)
mapOnType (forall {k} {k} (frep :: k) (trep :: k) (m :: * -> *).
Mapper frep trep m -> SubExp -> m SubExp
mapOnSubExp Mapper frep trep m
tv)) [Type]
ret
mapExp :: Mapper frep trep Identity -> Exp frep -> Exp trep
mapExp :: forall {k} {k} (frep :: k) (trep :: k).
Mapper frep trep Identity -> Exp frep -> Exp trep
mapExp Mapper frep trep Identity
m = forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} {k} (m :: * -> *) (frep :: k) (trep :: k).
Monad m =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper frep trep Identity
m
data Walker rep m = Walker
{ forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp :: SubExp -> m (),
forall {k} (rep :: k) (m :: * -> *).
Walker rep m -> Scope rep -> Body rep -> m ()
walkOnBody :: Scope rep -> Body rep -> m (),
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName :: VName -> m (),
forall {k} (rep :: k) (m :: * -> *).
Walker rep m -> RetType rep -> m ()
walkOnRetType :: RetType rep -> m (),
forall {k} (rep :: k) (m :: * -> *).
Walker rep m -> BranchType rep -> m ()
walkOnBranchType :: BranchType rep -> m (),
forall {k} (rep :: k) (m :: * -> *).
Walker rep m -> FParam rep -> m ()
walkOnFParam :: FParam rep -> m (),
forall {k} (rep :: k) (m :: * -> *).
Walker rep m -> LParam rep -> m ()
walkOnLParam :: LParam rep -> m (),
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> Op rep -> m ()
walkOnOp :: Op rep -> m ()
}
identityWalker :: Monad m => Walker rep m
identityWalker :: forall {k} (m :: * -> *) (rep :: k). Monad m => Walker rep m
identityWalker =
Walker
{ walkOnSubExp :: SubExp -> m ()
walkOnSubExp = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
walkOnBody :: Scope rep -> Body rep -> m ()
walkOnBody = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
walkOnVName :: VName -> m ()
walkOnVName = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
walkOnRetType :: RetType rep -> m ()
walkOnRetType = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
walkOnBranchType :: BranchType rep -> m ()
walkOnBranchType = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
walkOnFParam :: FParam rep -> m ()
walkOnFParam = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
walkOnLParam :: LParam rep -> m ()
walkOnLParam = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
walkOnOp :: Op rep -> m ()
walkOnOp = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
}
walkOnShape :: Monad m => Walker rep m -> Shape -> m ()
walkOnShape :: forall {k} (m :: * -> *) (rep :: k).
Monad m =>
Walker rep m -> Shape -> m ()
walkOnShape Walker rep m
tv (Shape [SubExp]
ds) = forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) [SubExp]
ds
walkOnType :: Monad m => Walker rep m -> Type -> m ()
walkOnType :: forall {k} (m :: * -> *) (rep :: k).
Monad m =>
Walker rep m -> Type -> m ()
walkOnType Walker rep m
_ Prim {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
walkOnType Walker rep m
tv (Acc VName
acc Shape
ispace [Type]
ts NoUniqueness
_) = do
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
acc
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) Shape
ispace
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (m :: * -> *) (rep :: k).
Monad m =>
Walker rep m -> Type -> m ()
walkOnType Walker rep m
tv) [Type]
ts
walkOnType Walker rep m
_ Mem {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
walkOnType Walker rep m
tv (Array PrimType
_ Shape
shape NoUniqueness
_) = forall {k} (m :: * -> *) (rep :: k).
Monad m =>
Walker rep m -> Shape -> m ()
walkOnShape Walker rep m
tv Shape
shape
walkOnLoopForm :: Monad m => Walker rep m -> LoopForm rep -> m ()
walkOnLoopForm :: forall {k} (m :: * -> *) (rep :: k).
Monad m =>
Walker rep m -> LoopForm rep -> m ()
walkOnLoopForm Walker rep m
tv (ForLoop VName
i IntType
_ SubExp
bound [(LParam rep, VName)]
loop_vars) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
i
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
bound
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *).
Walker rep m -> LParam rep -> m ()
walkOnLParam Walker rep m
tv) [LParam rep]
loop_lparams
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv) [VName]
loop_arrs
where
([LParam rep]
loop_lparams, [VName]
loop_arrs) = forall a b. [(a, b)] -> ([a], [b])
unzip [(LParam rep, VName)]
loop_vars
walkOnLoopForm Walker rep m
tv (WhileLoop VName
cond) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
cond
walkOnLambda :: Monad m => Walker rep m -> Lambda rep -> m ()
walkOnLambda :: forall {k} (m :: * -> *) (rep :: k).
Monad m =>
Walker rep m -> Lambda rep -> m ()
walkOnLambda Walker rep m
tv (Lambda [LParam rep]
params Body rep
body [Type]
ret) = do
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *).
Walker rep m -> LParam rep -> m ()
walkOnLParam Walker rep m
tv) [LParam rep]
params
forall {k} (rep :: k) (m :: * -> *).
Walker rep m -> Scope rep -> Body rep -> m ()
walkOnBody Walker rep m
tv (forall {k} (rep :: k) dec.
(LParamInfo rep ~ dec) =>
[Param dec] -> Scope rep
scopeOfLParams [LParam rep]
params) Body rep
body
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (m :: * -> *) (rep :: k).
Monad m =>
Walker rep m -> Type -> m ()
walkOnType Walker rep m
tv) [Type]
ret
walkExpM :: Monad m => Walker rep m -> Exp rep -> m ()
walkExpM :: forall {k} (m :: * -> *) (rep :: k).
Monad m =>
Walker rep m -> Exp rep -> m ()
walkExpM Walker rep m
tv (BasicOp (SubExp SubExp
se)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
se
walkExpM Walker rep m
tv (BasicOp (ArrayLit [SubExp]
els Type
rowt)) =
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) [SubExp]
els forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall {k} (m :: * -> *) (rep :: k).
Monad m =>
Walker rep m -> Type -> m ()
walkOnType Walker rep m
tv Type
rowt
walkExpM Walker rep m
tv (BasicOp (BinOp BinOp
_ SubExp
x SubExp
y)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
x forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
y
walkExpM Walker rep m
tv (BasicOp (CmpOp CmpOp
_ SubExp
x SubExp
y)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
x forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
y
walkExpM Walker rep m
tv (BasicOp (ConvOp ConvOp
_ SubExp
x)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
x
walkExpM Walker rep m
tv (BasicOp (UnOp UnOp
_ SubExp
x)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
x
walkExpM Walker rep m
tv (Match [SubExp]
ses [Case (Body rep)]
cases Body rep
defbody (MatchDec [BranchType rep]
ts MatchSort
_)) = do
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) [SubExp]
ses
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *).
Walker rep m -> Scope rep -> Body rep -> m ()
walkOnBody Walker rep m
tv forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall body. Case body -> body
caseBody) [Case (Body rep)]
cases
forall {k} (rep :: k) (m :: * -> *).
Walker rep m -> Scope rep -> Body rep -> m ()
walkOnBody Walker rep m
tv forall a. Monoid a => a
mempty Body rep
defbody
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *).
Walker rep m -> BranchType rep -> m ()
walkOnBranchType Walker rep m
tv) [BranchType rep]
ts
walkExpM Walker rep m
tv (Apply Name
_ [(SubExp, Diet)]
args [RetType rep]
ret (Safety, SrcLoc, [SrcLoc])
_) =
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(SubExp, Diet)]
args forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *).
Walker rep m -> RetType rep -> m ()
walkOnRetType Walker rep m
tv) [RetType rep]
ret
walkExpM Walker rep m
tv (BasicOp (Index VName
arr Slice SubExp
slice)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
arr forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) Slice SubExp
slice
walkExpM Walker rep m
tv (BasicOp (Update Safety
_ VName
arr Slice SubExp
slice SubExp
se)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
arr
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) Slice SubExp
slice
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
se
walkExpM Walker rep m
tv (BasicOp (FlatIndex VName
arr FlatSlice SubExp
slice)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
arr forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) FlatSlice SubExp
slice
walkExpM Walker rep m
tv (BasicOp (FlatUpdate VName
arr FlatSlice SubExp
slice VName
se)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
arr
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) FlatSlice SubExp
slice
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
se
walkExpM Walker rep m
tv (BasicOp (Iota SubExp
n SubExp
x SubExp
s IntType
_)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
n forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
x forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
s
walkExpM Walker rep m
tv (BasicOp (Replicate Shape
shape SubExp
vexp)) =
forall {k} (m :: * -> *) (rep :: k).
Monad m =>
Walker rep m -> Shape -> m ()
walkOnShape Walker rep m
tv Shape
shape forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
vexp
walkExpM Walker rep m
tv (BasicOp (Scratch PrimType
_ [SubExp]
shape)) =
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) [SubExp]
shape
walkExpM Walker rep m
tv (BasicOp (Reshape ReshapeKind
_ Shape
shape VName
arrexp)) =
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) Shape
shape forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
arrexp
walkExpM Walker rep m
tv (BasicOp (Rearrange [Int]
_ VName
e)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
e
walkExpM Walker rep m
tv (BasicOp (Rotate [SubExp]
es VName
e)) =
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) [SubExp]
es forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
e
walkExpM Walker rep m
tv (BasicOp (Concat Int
_ (VName
x :| [VName]
ys) SubExp
size)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
x forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv) [VName]
ys forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
size
walkExpM Walker rep m
tv (BasicOp (Copy VName
e)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
e
walkExpM Walker rep m
tv (BasicOp (Manifest [Int]
_ VName
e)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
e
walkExpM Walker rep m
tv (BasicOp (Assert SubExp
e ErrorMsg SubExp
msg (SrcLoc, [SrcLoc])
_)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
e forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) ErrorMsg SubExp
msg
walkExpM Walker rep m
tv (BasicOp (Opaque OpaqueOp
_ SubExp
e)) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv SubExp
e
walkExpM Walker rep m
tv (BasicOp (UpdateAcc VName
v [SubExp]
is [SubExp]
ses)) = do
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv VName
v
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) [SubExp]
is
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) [SubExp]
ses
walkExpM Walker rep m
tv (WithAcc [WithAccInput rep]
inputs Lambda rep
lam) = do
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [WithAccInput rep]
inputs forall a b. (a -> b) -> a -> b
$ \(Shape
shape, [VName]
vs, Maybe (Lambda rep, [SubExp])
op) -> do
forall {k} (m :: * -> *) (rep :: k).
Monad m =>
Walker rep m -> Shape -> m ()
walkOnShape Walker rep m
tv Shape
shape
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> VName -> m ()
walkOnVName Walker rep m
tv) [VName]
vs
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse (forall {k} (m :: * -> *) (rep :: k).
Monad m =>
Walker rep m -> Lambda rep -> m ()
walkOnLambda Walker rep m
tv) (forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv))) Maybe (Lambda rep, [SubExp])
op
forall {k} (m :: * -> *) (rep :: k).
Monad m =>
Walker rep m -> Lambda rep -> m ()
walkOnLambda Walker rep m
tv Lambda rep
lam
walkExpM Walker rep m
tv (DoLoop [(FParam rep, SubExp)]
merge LoopForm rep
form Body rep
loopbody) = do
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *).
Walker rep m -> FParam rep -> m ()
walkOnFParam Walker rep m
tv) [FParam rep]
params
forall {k} (m :: * -> *) (rep :: k).
Monad m =>
Walker rep m -> LoopForm rep -> m ()
walkOnLoopForm Walker rep m
tv LoopForm rep
form
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall {k} (rep :: k) (m :: * -> *). Walker rep m -> SubExp -> m ()
walkOnSubExp Walker rep m
tv) [SubExp]
args
let scope :: Scope rep
scope = forall {k} (rep :: k) dec.
(FParamInfo rep ~ dec) =>
[Param dec] -> Scope rep
scopeOfFParams [FParam rep]
params forall a. Semigroup a => a -> a -> a
<> forall {k} (rep :: k) a. Scoped rep a => a -> Scope rep
scopeOf LoopForm rep
form
forall {k} (rep :: k) (m :: * -> *).
Walker rep m -> Scope rep -> Body rep -> m ()
walkOnBody Walker rep m
tv Scope rep
scope Body rep
loopbody
where
([FParam rep]
params, [SubExp]
args) = forall a b. [(a, b)] -> ([a], [b])
unzip [(FParam rep, SubExp)]
merge
walkExpM Walker rep m
tv (Op Op rep
op) =
forall {k} (rep :: k) (m :: * -> *). Walker rep m -> Op rep -> m ()
walkOnOp Walker rep m
tv Op rep
op
type OpStmsTraverser m op rep = (Scope rep -> Stms rep -> m (Stms rep)) -> op -> m op
class TraverseOpStms rep where
traverseOpStms :: Monad m => OpStmsTraverser m (Op rep) rep
traverseLambdaStms :: Monad m => OpStmsTraverser m (Lambda rep) rep
traverseLambdaStms :: forall {k} (m :: * -> *) (rep :: k).
Monad m =>
OpStmsTraverser m (Lambda rep) rep
traverseLambdaStms Scope rep -> Stms rep -> m (Stms rep)
f (Lambda [LParam rep]
ps (Body BodyDec rep
dec Stms rep
stms Result
res) [Type]
ret) =
forall {k} (rep :: k).
[LParam rep] -> Body rep -> [Type] -> Lambda rep
Lambda [LParam rep]
ps forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall {k} (rep :: k).
BodyDec rep -> Stms rep -> Result -> Body rep
Body BodyDec rep
dec forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Scope rep -> Stms rep -> m (Stms rep)
f (forall {k} (rep :: k) dec.
(LParamInfo rep ~ dec) =>
[Param dec] -> Scope rep
scopeOfLParams [LParam rep]
ps) Stms rep
stms forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure Result
res) forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure [Type]
ret