{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
module Futhark.IR.MC.Op
( MCOp (..),
traverseMCOpStms,
typeCheckMCOp,
simplifyMCOp,
module Futhark.IR.SegOp,
)
where
import Data.Bifunctor (first)
import Futhark.Analysis.Metrics
import Futhark.Analysis.SymbolTable qualified as ST
import Futhark.IR
import Futhark.IR.Aliases (Aliases, CanBeAliased (..))
import Futhark.IR.Prop.Aliases
import Futhark.IR.SegOp
import Futhark.IR.TypeCheck qualified as TC
import Futhark.Optimise.Simplify qualified as Simplify
import Futhark.Optimise.Simplify.Engine qualified as Engine
import Futhark.Optimise.Simplify.Rep
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
import Futhark.Util.Pretty
( nestedBlock,
pretty,
(<+>),
(</>),
)
import Prelude hiding (id, (.))
data MCOp op rep
=
ParOp
(Maybe (SegOp () rep))
(SegOp () rep)
|
OtherOp (op rep)
deriving (MCOp op rep -> MCOp op rep -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall (op :: * -> *) rep.
(RepTypes rep, Eq (op rep)) =>
MCOp op rep -> MCOp op rep -> Bool
/= :: MCOp op rep -> MCOp op rep -> Bool
$c/= :: forall (op :: * -> *) rep.
(RepTypes rep, Eq (op rep)) =>
MCOp op rep -> MCOp op rep -> Bool
== :: MCOp op rep -> MCOp op rep -> Bool
$c== :: forall (op :: * -> *) rep.
(RepTypes rep, Eq (op rep)) =>
MCOp op rep -> MCOp op rep -> Bool
Eq, MCOp op rep -> MCOp op rep -> Bool
MCOp op rep -> MCOp op rep -> Ordering
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall {op :: * -> *} {rep}.
(RepTypes rep, Ord (op rep)) =>
Eq (MCOp op rep)
forall (op :: * -> *) rep.
(RepTypes rep, Ord (op rep)) =>
MCOp op rep -> MCOp op rep -> Bool
forall (op :: * -> *) rep.
(RepTypes rep, Ord (op rep)) =>
MCOp op rep -> MCOp op rep -> Ordering
forall (op :: * -> *) rep.
(RepTypes rep, Ord (op rep)) =>
MCOp op rep -> MCOp op rep -> MCOp op rep
min :: MCOp op rep -> MCOp op rep -> MCOp op rep
$cmin :: forall (op :: * -> *) rep.
(RepTypes rep, Ord (op rep)) =>
MCOp op rep -> MCOp op rep -> MCOp op rep
max :: MCOp op rep -> MCOp op rep -> MCOp op rep
$cmax :: forall (op :: * -> *) rep.
(RepTypes rep, Ord (op rep)) =>
MCOp op rep -> MCOp op rep -> MCOp op rep
>= :: MCOp op rep -> MCOp op rep -> Bool
$c>= :: forall (op :: * -> *) rep.
(RepTypes rep, Ord (op rep)) =>
MCOp op rep -> MCOp op rep -> Bool
> :: MCOp op rep -> MCOp op rep -> Bool
$c> :: forall (op :: * -> *) rep.
(RepTypes rep, Ord (op rep)) =>
MCOp op rep -> MCOp op rep -> Bool
<= :: MCOp op rep -> MCOp op rep -> Bool
$c<= :: forall (op :: * -> *) rep.
(RepTypes rep, Ord (op rep)) =>
MCOp op rep -> MCOp op rep -> Bool
< :: MCOp op rep -> MCOp op rep -> Bool
$c< :: forall (op :: * -> *) rep.
(RepTypes rep, Ord (op rep)) =>
MCOp op rep -> MCOp op rep -> Bool
compare :: MCOp op rep -> MCOp op rep -> Ordering
$ccompare :: forall (op :: * -> *) rep.
(RepTypes rep, Ord (op rep)) =>
MCOp op rep -> MCOp op rep -> Ordering
Ord, Int -> MCOp op rep -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall (op :: * -> *) rep.
(RepTypes rep, Show (op rep)) =>
Int -> MCOp op rep -> ShowS
forall (op :: * -> *) rep.
(RepTypes rep, Show (op rep)) =>
[MCOp op rep] -> ShowS
forall (op :: * -> *) rep.
(RepTypes rep, Show (op rep)) =>
MCOp op rep -> String
showList :: [MCOp op rep] -> ShowS
$cshowList :: forall (op :: * -> *) rep.
(RepTypes rep, Show (op rep)) =>
[MCOp op rep] -> ShowS
show :: MCOp op rep -> String
$cshow :: forall (op :: * -> *) rep.
(RepTypes rep, Show (op rep)) =>
MCOp op rep -> String
showsPrec :: Int -> MCOp op rep -> ShowS
$cshowsPrec :: forall (op :: * -> *) rep.
(RepTypes rep, Show (op rep)) =>
Int -> MCOp op rep -> ShowS
Show)
traverseMCOpStms ::
Monad m =>
OpStmsTraverser m (op rep) rep ->
OpStmsTraverser m (MCOp op rep) rep
traverseMCOpStms :: forall (m :: * -> *) (op :: * -> *) rep.
Monad m =>
OpStmsTraverser m (op rep) rep
-> OpStmsTraverser m (MCOp op rep) rep
traverseMCOpStms OpStmsTraverser m (op rep) rep
_ Scope rep -> Stms rep -> m (Stms rep)
f (ParOp Maybe (SegOp () rep)
par_op SegOp () rep
op) =
forall (op :: * -> *) rep.
Maybe (SegOp () rep) -> SegOp () rep -> MCOp op rep
ParOp forall (f :: * -> *) a b. Functor 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 (m :: * -> *) lvl rep.
Monad m =>
OpStmsTraverser m (SegOp lvl rep) rep
traverseSegOpStms Scope rep -> Stms rep -> m (Stms rep)
f) Maybe (SegOp () rep)
par_op forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (m :: * -> *) lvl rep.
Monad m =>
OpStmsTraverser m (SegOp lvl rep) rep
traverseSegOpStms Scope rep -> Stms rep -> m (Stms rep)
f SegOp () rep
op
traverseMCOpStms OpStmsTraverser m (op rep) rep
onInner Scope rep -> Stms rep -> m (Stms rep)
f (OtherOp op rep
op) = forall (op :: * -> *) rep. op rep -> MCOp op rep
OtherOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpStmsTraverser m (op rep) rep
onInner Scope rep -> Stms rep -> m (Stms rep)
f op rep
op
instance (ASTRep rep, Substitute (op rep)) => Substitute (MCOp op rep) where
substituteNames :: Map VName VName -> MCOp op rep -> MCOp op rep
substituteNames Map VName VName
substs (ParOp Maybe (SegOp () rep)
par_op SegOp () rep
op) =
forall (op :: * -> *) rep.
Maybe (SegOp () rep) -> SegOp () rep -> MCOp op rep
ParOp (forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (SegOp () rep)
par_op) (forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs SegOp () rep
op)
substituteNames Map VName VName
substs (OtherOp op rep
op) =
forall (op :: * -> *) rep. op rep -> MCOp op rep
OtherOp forall a b. (a -> b) -> a -> b
$ forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs op rep
op
instance (ASTRep rep, Rename (op rep)) => Rename (MCOp op rep) where
rename :: MCOp op rep -> RenameM (MCOp op rep)
rename (ParOp Maybe (SegOp () rep)
par_op SegOp () rep
op) = forall (op :: * -> *) rep.
Maybe (SegOp () rep) -> SegOp () rep -> MCOp op rep
ParOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Rename a => a -> RenameM a
rename Maybe (SegOp () rep)
par_op forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Rename a => a -> RenameM a
rename SegOp () rep
op
rename (OtherOp op rep
op) = forall (op :: * -> *) rep. op rep -> MCOp op rep
OtherOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Rename a => a -> RenameM a
rename op rep
op
instance (ASTRep rep, FreeIn (op rep)) => FreeIn (MCOp op rep) where
freeIn' :: MCOp op rep -> FV
freeIn' (ParOp Maybe (SegOp () rep)
par_op SegOp () rep
op) = forall a. FreeIn a => a -> FV
freeIn' Maybe (SegOp () rep)
par_op forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' SegOp () rep
op
freeIn' (OtherOp op rep
op) = forall a. FreeIn a => a -> FV
freeIn' op rep
op
instance (ASTRep rep, IsOp (op rep)) => IsOp (MCOp op rep) where
safeOp :: MCOp op rep -> Bool
safeOp (ParOp Maybe (SegOp () rep)
_ SegOp () rep
op) = forall op. IsOp op => op -> Bool
safeOp SegOp () rep
op
safeOp (OtherOp op rep
op) = forall op. IsOp op => op -> Bool
safeOp op rep
op
cheapOp :: MCOp op rep -> Bool
cheapOp (ParOp Maybe (SegOp () rep)
_ SegOp () rep
op) = forall op. IsOp op => op -> Bool
cheapOp SegOp () rep
op
cheapOp (OtherOp op rep
op) = forall op. IsOp op => op -> Bool
cheapOp op rep
op
instance TypedOp (op rep) => TypedOp (MCOp op rep) where
opType :: forall t (m :: * -> *). HasScope t m => MCOp op rep -> m [ExtType]
opType (ParOp Maybe (SegOp () rep)
_ SegOp () rep
op) = forall op t (m :: * -> *).
(TypedOp op, HasScope t m) =>
op -> m [ExtType]
opType SegOp () rep
op
opType (OtherOp op rep
op) = forall op t (m :: * -> *).
(TypedOp op, HasScope t m) =>
op -> m [ExtType]
opType op rep
op
instance (Aliased rep, AliasedOp (op rep)) => AliasedOp (MCOp op rep) where
opAliases :: MCOp op rep -> [Names]
opAliases (ParOp Maybe (SegOp () rep)
_ SegOp () rep
op) = forall op. AliasedOp op => op -> [Names]
opAliases SegOp () rep
op
opAliases (OtherOp op rep
op) = forall op. AliasedOp op => op -> [Names]
opAliases op rep
op
consumedInOp :: MCOp op rep -> Names
consumedInOp (ParOp Maybe (SegOp () rep)
_ SegOp () rep
op) = forall op. AliasedOp op => op -> Names
consumedInOp SegOp () rep
op
consumedInOp (OtherOp op rep
op) = forall op. AliasedOp op => op -> Names
consumedInOp op rep
op
instance CanBeAliased op => CanBeAliased (MCOp op) where
addOpAliases :: forall rep.
AliasableRep rep =>
AliasTable -> MCOp op rep -> MCOp op (Aliases rep)
addOpAliases AliasTable
aliases (ParOp Maybe (SegOp () rep)
par_op SegOp () rep
op) =
forall (op :: * -> *) rep.
Maybe (SegOp () rep) -> SegOp () rep -> MCOp op rep
ParOp (forall (op :: * -> *) rep.
(CanBeAliased op, AliasableRep rep) =>
AliasTable -> op rep -> op (Aliases rep)
addOpAliases AliasTable
aliases forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (SegOp () rep)
par_op) (forall (op :: * -> *) rep.
(CanBeAliased op, AliasableRep rep) =>
AliasTable -> op rep -> op (Aliases rep)
addOpAliases AliasTable
aliases SegOp () rep
op)
addOpAliases AliasTable
aliases (OtherOp op rep
op) =
forall (op :: * -> *) rep. op rep -> MCOp op rep
OtherOp forall a b. (a -> b) -> a -> b
$ forall (op :: * -> *) rep.
(CanBeAliased op, AliasableRep rep) =>
AliasTable -> op rep -> op (Aliases rep)
addOpAliases AliasTable
aliases op rep
op
instance CanBeWise op => CanBeWise (MCOp op) where
addOpWisdom :: forall rep. Informing rep => MCOp op rep -> MCOp op (Wise rep)
addOpWisdom (ParOp Maybe (SegOp () rep)
par_op SegOp () rep
op) =
forall (op :: * -> *) rep.
Maybe (SegOp () rep) -> SegOp () rep -> MCOp op rep
ParOp (forall (op :: * -> *) rep.
(CanBeWise op, Informing rep) =>
op rep -> op (Wise rep)
addOpWisdom forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (SegOp () rep)
par_op) (forall (op :: * -> *) rep.
(CanBeWise op, Informing rep) =>
op rep -> op (Wise rep)
addOpWisdom SegOp () rep
op)
addOpWisdom (OtherOp op rep
op) =
forall (op :: * -> *) rep. op rep -> MCOp op rep
OtherOp forall a b. (a -> b) -> a -> b
$ forall (op :: * -> *) rep.
(CanBeWise op, Informing rep) =>
op rep -> op (Wise rep)
addOpWisdom op rep
op
instance (ASTRep rep, ST.IndexOp (op rep)) => ST.IndexOp (MCOp op rep) where
indexOp :: forall rep.
(ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep
-> Int -> MCOp op rep -> [TPrimExp Int64 VName] -> Maybe Indexed
indexOp SymbolTable rep
vtable Int
k (ParOp Maybe (SegOp () rep)
_ SegOp () rep
op) [TPrimExp Int64 VName]
is = forall op rep.
(IndexOp op, ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep
-> Int -> op -> [TPrimExp Int64 VName] -> Maybe Indexed
ST.indexOp SymbolTable rep
vtable Int
k SegOp () rep
op [TPrimExp Int64 VName]
is
indexOp SymbolTable rep
vtable Int
k (OtherOp op rep
op) [TPrimExp Int64 VName]
is = forall op rep.
(IndexOp op, ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep
-> Int -> op -> [TPrimExp Int64 VName] -> Maybe Indexed
ST.indexOp SymbolTable rep
vtable Int
k op rep
op [TPrimExp Int64 VName]
is
instance (PrettyRep rep, Pretty (op rep)) => Pretty (MCOp op rep) where
pretty :: forall ann. MCOp op rep -> Doc ann
pretty (ParOp Maybe (SegOp () rep)
Nothing SegOp () rep
op) = forall a ann. Pretty a => a -> Doc ann
pretty SegOp () rep
op
pretty (ParOp (Just SegOp () rep
par_op) SegOp () rep
op) =
Doc ann
"par"
forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall a. Doc a -> Doc a -> Doc a -> Doc a
nestedBlock Doc ann
"{" Doc ann
"}" (forall a ann. Pretty a => a -> Doc ann
pretty SegOp () rep
par_op)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc ann
"seq"
forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall a. Doc a -> Doc a -> Doc a -> Doc a
nestedBlock Doc ann
"{" Doc ann
"}" (forall a ann. Pretty a => a -> Doc ann
pretty SegOp () rep
op)
pretty (OtherOp op rep
op) = forall a ann. Pretty a => a -> Doc ann
pretty op rep
op
instance (OpMetrics (Op rep), OpMetrics (op rep)) => OpMetrics (MCOp op rep) where
opMetrics :: MCOp op rep -> MetricsM ()
opMetrics (ParOp Maybe (SegOp () rep)
par_op SegOp () rep
op) = forall op. OpMetrics op => op -> MetricsM ()
opMetrics Maybe (SegOp () rep)
par_op forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall op. OpMetrics op => op -> MetricsM ()
opMetrics SegOp () rep
op
opMetrics (OtherOp op rep
op) = forall op. OpMetrics op => op -> MetricsM ()
opMetrics op rep
op
instance RephraseOp op => RephraseOp (MCOp op) where
rephraseInOp :: forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> MCOp op from -> m (MCOp op to)
rephraseInOp Rephraser m from to
r (ParOp Maybe (SegOp () from)
par_op SegOp () from
op) =
forall (op :: * -> *) rep.
Maybe (SegOp () rep) -> SegOp () rep -> MCOp op rep
ParOp forall (f :: * -> *) a b. Functor 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 (op :: * -> *) (m :: * -> *) from to.
(RephraseOp op, Monad m) =>
Rephraser m from to -> op from -> m (op to)
rephraseInOp Rephraser m from to
r) Maybe (SegOp () from)
par_op forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (op :: * -> *) (m :: * -> *) from to.
(RephraseOp op, Monad m) =>
Rephraser m from to -> op from -> m (op to)
rephraseInOp Rephraser m from to
r SegOp () from
op
rephraseInOp Rephraser m from to
r (OtherOp op from
op) = forall (op :: * -> *) rep. op rep -> MCOp op rep
OtherOp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (op :: * -> *) (m :: * -> *) from to.
(RephraseOp op, Monad m) =>
Rephraser m from to -> op from -> m (op to)
rephraseInOp Rephraser m from to
r op from
op
typeCheckMCOp ::
TC.Checkable rep =>
(op (Aliases rep) -> TC.TypeM rep ()) ->
MCOp op (Aliases rep) ->
TC.TypeM rep ()
typeCheckMCOp :: forall rep (op :: * -> *).
Checkable rep =>
(op (Aliases rep) -> TypeM rep ())
-> MCOp op (Aliases rep) -> TypeM rep ()
typeCheckMCOp op (Aliases rep) -> TypeM rep ()
_ (ParOp (Just SegOp () (Aliases rep)
par_op) SegOp () (Aliases rep)
op) = do
((), ())
_ <- forall rep lvl.
Checkable rep =>
(lvl -> TypeM rep ()) -> SegOp lvl (Aliases rep) -> TypeM rep ()
typeCheckSegOp forall (f :: * -> *) a. Applicative f => a -> f a
pure SegOp () (Aliases rep)
par_op forall rep a b. TypeM rep a -> TypeM rep b -> TypeM rep (a, b)
`TC.alternative` forall rep lvl.
Checkable rep =>
(lvl -> TypeM rep ()) -> SegOp lvl (Aliases rep) -> TypeM rep ()
typeCheckSegOp forall (f :: * -> *) a. Applicative f => a -> f a
pure SegOp () (Aliases rep)
op
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
typeCheckMCOp op (Aliases rep) -> TypeM rep ()
_ (ParOp Maybe (SegOp () (Aliases rep))
Nothing SegOp () (Aliases rep)
op) =
forall rep lvl.
Checkable rep =>
(lvl -> TypeM rep ()) -> SegOp lvl (Aliases rep) -> TypeM rep ()
typeCheckSegOp forall (f :: * -> *) a. Applicative f => a -> f a
pure SegOp () (Aliases rep)
op
typeCheckMCOp op (Aliases rep) -> TypeM rep ()
f (OtherOp op (Aliases rep)
op) = op (Aliases rep) -> TypeM rep ()
f op (Aliases rep)
op
simplifyMCOp ::
( Engine.SimplifiableRep rep,
BodyDec rep ~ ()
) =>
Simplify.SimplifyOp rep (op (Wise rep)) ->
MCOp op (Wise rep) ->
Engine.SimpleM rep (MCOp op (Wise rep), Stms (Wise rep))
simplifyMCOp :: forall rep (op :: * -> *).
(SimplifiableRep rep, BodyDec rep ~ ()) =>
SimplifyOp rep (op (Wise rep))
-> MCOp op (Wise rep)
-> SimpleM rep (MCOp op (Wise rep), Stms (Wise rep))
simplifyMCOp SimplifyOp rep (op (Wise rep))
f (OtherOp op (Wise rep)
op) = do
(op (Wise rep)
op', Stms (Wise rep)
stms) <- SimplifyOp rep (op (Wise rep))
f op (Wise rep)
op
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (op :: * -> *) rep. op rep -> MCOp op rep
OtherOp op (Wise rep)
op', Stms (Wise rep)
stms)
simplifyMCOp SimplifyOp rep (op (Wise rep))
_ (ParOp Maybe (SegOp () (Wise rep))
par_op SegOp () (Wise rep)
op) = do
(Maybe (SegOp () (Wise rep))
par_op', Stms (Wise rep)
par_op_hoisted) <-
case Maybe (SegOp () (Wise rep))
par_op of
Maybe (SegOp () (Wise rep))
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Maybe a
Nothing, forall a. Monoid a => a
mempty)
Just SegOp () (Wise rep)
x -> forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep lvl.
(SimplifiableRep rep, BodyDec rep ~ (), Simplifiable lvl) =>
SegOp lvl (Wise rep)
-> SimpleM rep (SegOp lvl (Wise rep), Stms (Wise rep))
simplifySegOp SegOp () (Wise rep)
x
(SegOp () (Wise rep)
op', Stms (Wise rep)
op_hoisted) <- forall rep lvl.
(SimplifiableRep rep, BodyDec rep ~ (), Simplifiable lvl) =>
SegOp lvl (Wise rep)
-> SimpleM rep (SegOp lvl (Wise rep), Stms (Wise rep))
simplifySegOp SegOp () (Wise rep)
op
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (op :: * -> *) rep.
Maybe (SegOp () rep) -> SegOp () rep -> MCOp op rep
ParOp Maybe (SegOp () (Wise rep))
par_op' SegOp () (Wise rep)
op', Stms (Wise rep)
par_op_hoisted forall a. Semigroup a => a -> a -> a
<> Stms (Wise rep)
op_hoisted)