{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- | Definitions for multicore operations.
--
-- Most of the interesting stuff is in "Futhark.IR.SegOp", which is
-- also re-exported from here.
module Futhark.IR.MC.Op
  ( MCOp (..),
    typeCheckMCOp,
    simplifyMCOp,
    module Futhark.IR.SegOp,
  )
where

import Control.Category
import Data.Bifunctor (first)
import Futhark.Analysis.Metrics
import qualified Futhark.Analysis.SymbolTable as ST
import Futhark.IR
import Futhark.IR.Aliases (Aliases)
import Futhark.IR.Prop.Aliases
import Futhark.IR.SegOp
import qualified Futhark.Optimise.Simplify as Simplify
import qualified Futhark.Optimise.Simplify.Engine as Engine
import Futhark.Optimise.Simplify.Lore
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
import qualified Futhark.TypeCheck as TC
import Futhark.Util.Pretty
  ( Pretty,
    nestedBlock,
    ppr,
    (<+>),
    (</>),
  )
import GHC.Generics (Generic)
import Language.SexpGrammar as Sexp
import Language.SexpGrammar.Generic
import Prelude hiding (id, (.))

-- | An operation for the multicore representation.  Feel free to
-- extend this on an ad hoc basis as needed.  Parameterised with some
-- other operation.
data MCOp lore op
  = -- | The first 'SegOp' (if it exists) contains nested parallelism,
    -- while the second one has a fully sequential body.  They are
    -- semantically fully equivalent.
    ParOp
      (Maybe (SegOp () lore))
      (SegOp () lore)
  | -- | Something else (in practice often a SOAC).
    OtherOp op
  deriving (MCOp lore op -> MCOp lore op -> Bool
(MCOp lore op -> MCOp lore op -> Bool)
-> (MCOp lore op -> MCOp lore op -> Bool) -> Eq (MCOp lore op)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall lore op.
(Decorations lore, Eq op) =>
MCOp lore op -> MCOp lore op -> Bool
/= :: MCOp lore op -> MCOp lore op -> Bool
$c/= :: forall lore op.
(Decorations lore, Eq op) =>
MCOp lore op -> MCOp lore op -> Bool
== :: MCOp lore op -> MCOp lore op -> Bool
$c== :: forall lore op.
(Decorations lore, Eq op) =>
MCOp lore op -> MCOp lore op -> Bool
Eq, Eq (MCOp lore op)
Eq (MCOp lore op)
-> (MCOp lore op -> MCOp lore op -> Ordering)
-> (MCOp lore op -> MCOp lore op -> Bool)
-> (MCOp lore op -> MCOp lore op -> Bool)
-> (MCOp lore op -> MCOp lore op -> Bool)
-> (MCOp lore op -> MCOp lore op -> Bool)
-> (MCOp lore op -> MCOp lore op -> MCOp lore op)
-> (MCOp lore op -> MCOp lore op -> MCOp lore op)
-> Ord (MCOp lore op)
MCOp lore op -> MCOp lore op -> Bool
MCOp lore op -> MCOp lore op -> Ordering
MCOp lore op -> MCOp lore op -> MCOp lore op
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 lore op. (Decorations lore, Ord op) => Eq (MCOp lore op)
forall lore op.
(Decorations lore, Ord op) =>
MCOp lore op -> MCOp lore op -> Bool
forall lore op.
(Decorations lore, Ord op) =>
MCOp lore op -> MCOp lore op -> Ordering
forall lore op.
(Decorations lore, Ord op) =>
MCOp lore op -> MCOp lore op -> MCOp lore op
min :: MCOp lore op -> MCOp lore op -> MCOp lore op
$cmin :: forall lore op.
(Decorations lore, Ord op) =>
MCOp lore op -> MCOp lore op -> MCOp lore op
max :: MCOp lore op -> MCOp lore op -> MCOp lore op
$cmax :: forall lore op.
(Decorations lore, Ord op) =>
MCOp lore op -> MCOp lore op -> MCOp lore op
>= :: MCOp lore op -> MCOp lore op -> Bool
$c>= :: forall lore op.
(Decorations lore, Ord op) =>
MCOp lore op -> MCOp lore op -> Bool
> :: MCOp lore op -> MCOp lore op -> Bool
$c> :: forall lore op.
(Decorations lore, Ord op) =>
MCOp lore op -> MCOp lore op -> Bool
<= :: MCOp lore op -> MCOp lore op -> Bool
$c<= :: forall lore op.
(Decorations lore, Ord op) =>
MCOp lore op -> MCOp lore op -> Bool
< :: MCOp lore op -> MCOp lore op -> Bool
$c< :: forall lore op.
(Decorations lore, Ord op) =>
MCOp lore op -> MCOp lore op -> Bool
compare :: MCOp lore op -> MCOp lore op -> Ordering
$ccompare :: forall lore op.
(Decorations lore, Ord op) =>
MCOp lore op -> MCOp lore op -> Ordering
$cp1Ord :: forall lore op. (Decorations lore, Ord op) => Eq (MCOp lore op)
Ord, Int -> MCOp lore op -> ShowS
[MCOp lore op] -> ShowS
MCOp lore op -> String
(Int -> MCOp lore op -> ShowS)
-> (MCOp lore op -> String)
-> ([MCOp lore op] -> ShowS)
-> Show (MCOp lore op)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall lore op.
(Decorations lore, Show op) =>
Int -> MCOp lore op -> ShowS
forall lore op.
(Decorations lore, Show op) =>
[MCOp lore op] -> ShowS
forall lore op.
(Decorations lore, Show op) =>
MCOp lore op -> String
showList :: [MCOp lore op] -> ShowS
$cshowList :: forall lore op.
(Decorations lore, Show op) =>
[MCOp lore op] -> ShowS
show :: MCOp lore op -> String
$cshow :: forall lore op.
(Decorations lore, Show op) =>
MCOp lore op -> String
showsPrec :: Int -> MCOp lore op -> ShowS
$cshowsPrec :: forall lore op.
(Decorations lore, Show op) =>
Int -> MCOp lore op -> ShowS
Show, (forall x. MCOp lore op -> Rep (MCOp lore op) x)
-> (forall x. Rep (MCOp lore op) x -> MCOp lore op)
-> Generic (MCOp lore op)
forall x. Rep (MCOp lore op) x -> MCOp lore op
forall x. MCOp lore op -> Rep (MCOp lore op) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall lore op x. Rep (MCOp lore op) x -> MCOp lore op
forall lore op x. MCOp lore op -> Rep (MCOp lore op) x
$cto :: forall lore op x. Rep (MCOp lore op) x -> MCOp lore op
$cfrom :: forall lore op x. MCOp lore op -> Rep (MCOp lore op) x
Generic)

instance (Decorations lore, SexpIso op) => SexpIso (MCOp lore op) where
  sexpIso :: Grammar Position (Sexp :- t) (MCOp lore op :- t)
sexpIso =
    Coproduct
  Position
  (Sexp :- t)
  '[SegOp () lore :- (Maybe (SegOp () lore) :- t), op :- t]
  (MCOp lore op)
  t
-> Grammar Position (Sexp :- t) (MCOp lore op :- t)
forall a (bs :: [*]) t p s.
(Generic a, MkPrismList (Rep a), Match (Rep a) bs t,
 bs ~ Coll (Rep a) t) =>
Coproduct p s bs a t -> Grammar p s (a :- t)
match (Coproduct
   Position
   (Sexp :- t)
   '[SegOp () lore :- (Maybe (SegOp () lore) :- t), op :- t]
   (MCOp lore op)
   t
 -> Grammar Position (Sexp :- t) (MCOp lore op :- t))
-> Coproduct
     Position
     (Sexp :- t)
     '[SegOp () lore :- (Maybe (SegOp () lore) :- t), op :- t]
     (MCOp lore op)
     t
-> Grammar Position (Sexp :- t) (MCOp lore op :- t)
forall a b. (a -> b) -> a -> b
$
      (Grammar
   Position
   (SegOp () lore :- (Maybe (SegOp () lore) :- t))
   (MCOp lore op :- t)
 -> Grammar Position (Sexp :- t) (MCOp lore op :- t))
-> Coproduct Position (Sexp :- t) '[op :- t] (MCOp lore op) t
-> Coproduct
     Position
     (Sexp :- t)
     '[SegOp () lore :- (Maybe (SegOp () lore) :- t), op :- t]
     (MCOp lore op)
     t
forall p b a t s (bs1 :: [*]).
(Grammar p b (a :- t) -> Grammar p s (a :- t))
-> Coproduct p s bs1 a t -> Coproduct p s (b : bs1) a t
With (Grammar
  Position
  (SegOp () lore :- (Maybe (SegOp () lore) :- t))
  (MCOp lore op :- t)
-> Grammar
     Position
     (Sexp :- t)
     (SegOp () lore :- (Maybe (SegOp () lore) :- t))
-> Grammar Position (Sexp :- t) (MCOp lore op :- t)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Grammar
  Position
  (List :- t)
  (List :- (SegOp () lore :- (Maybe (SegOp () lore) :- t)))
-> Grammar
     Position
     (Sexp :- t)
     (SegOp () lore :- (Maybe (SegOp () lore) :- t))
forall t t'.
Grammar Position (List :- t) (List :- t')
-> Grammar Position (Sexp :- t) t'
Sexp.list (Grammar Position (Sexp :- t) (Maybe (SegOp () lore) :- t)
-> Grammar
     Position (List :- t) (List :- (Maybe (SegOp () lore) :- t))
forall t t'.
Grammar Position (Sexp :- t) t'
-> Grammar Position (List :- t) (List :- t')
Sexp.el Grammar Position (Sexp :- t) (Maybe (SegOp () lore) :- t)
forall a. SexpIso a => SexpGrammar a
sexpIso Grammar Position (List :- t) (List :- (Maybe (SegOp () lore) :- t))
-> Grammar
     Position
     (List :- (Maybe (SegOp () lore) :- t))
     (List :- (SegOp () lore :- (Maybe (SegOp () lore) :- t)))
-> Grammar
     Position
     (List :- t)
     (List :- (SegOp () lore :- (Maybe (SegOp () lore) :- t)))
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Grammar
  Position
  (Sexp :- (Maybe (SegOp () lore) :- t))
  (SegOp () lore :- (Maybe (SegOp () lore) :- t))
-> Grammar
     Position
     (List :- (Maybe (SegOp () lore) :- t))
     (List :- (SegOp () lore :- (Maybe (SegOp () lore) :- t)))
forall t t'.
Grammar Position (Sexp :- t) t'
-> Grammar Position (List :- t) (List :- t')
Sexp.el Grammar
  Position
  (Sexp :- (Maybe (SegOp () lore) :- t))
  (SegOp () lore :- (Maybe (SegOp () lore) :- t))
forall a. SexpIso a => SexpGrammar a
sexpIso)) (Coproduct Position (Sexp :- t) '[op :- t] (MCOp lore op) t
 -> Coproduct
      Position
      (Sexp :- t)
      '[SegOp () lore :- (Maybe (SegOp () lore) :- t), op :- t]
      (MCOp lore op)
      t)
-> Coproduct Position (Sexp :- t) '[op :- t] (MCOp lore op) t
-> Coproduct
     Position
     (Sexp :- t)
     '[SegOp () lore :- (Maybe (SegOp () lore) :- t), op :- t]
     (MCOp lore op)
     t
forall a b. (a -> b) -> a -> b
$
        (Grammar Position (op :- t) (MCOp lore op :- t)
 -> Grammar Position (Sexp :- t) (MCOp lore op :- t))
-> Coproduct Position (Sexp :- t) '[] (MCOp lore op) t
-> Coproduct Position (Sexp :- t) '[op :- t] (MCOp lore op) t
forall p b a t s (bs1 :: [*]).
(Grammar p b (a :- t) -> Grammar p s (a :- t))
-> Coproduct p s bs1 a t -> Coproduct p s (b : bs1) a t
With
          (Grammar Position (op :- t) (MCOp lore op :- t)
-> Grammar Position (Sexp :- t) (op :- t)
-> Grammar Position (Sexp :- t) (MCOp lore op :- t)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Grammar Position (List :- t) (List :- (op :- t))
-> Grammar Position (Sexp :- t) (op :- t)
forall t t'.
Grammar Position (List :- t) (List :- t')
-> Grammar Position (Sexp :- t) t'
Sexp.list (Grammar Position (Sexp :- t) (op :- t)
-> Grammar Position (List :- t) (List :- (op :- t))
forall t t'.
Grammar Position (Sexp :- t) t'
-> Grammar Position (List :- t) (List :- t')
Sexp.el Grammar Position (Sexp :- t) (op :- t)
forall a. SexpIso a => SexpGrammar a
sexpIso))
          Coproduct Position (Sexp :- t) '[] (MCOp lore op) t
forall p s a t. Coproduct p s '[] a t
End

instance (ASTLore lore, Substitute op) => Substitute (MCOp lore op) where
  substituteNames :: Map VName VName -> MCOp lore op -> MCOp lore op
substituteNames Map VName VName
substs (ParOp Maybe (SegOp () lore)
par_op SegOp () lore
op) =
    Maybe (SegOp () lore) -> SegOp () lore -> MCOp lore op
forall lore op.
Maybe (SegOp () lore) -> SegOp () lore -> MCOp lore op
ParOp (Map VName VName -> SegOp () lore -> SegOp () lore
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs (SegOp () lore -> SegOp () lore)
-> Maybe (SegOp () lore) -> Maybe (SegOp () lore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (SegOp () lore)
par_op) (Map VName VName -> SegOp () lore -> SegOp () lore
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs SegOp () lore
op)
  substituteNames Map VName VName
substs (OtherOp op
op) =
    op -> MCOp lore op
forall lore op. op -> MCOp lore op
OtherOp (op -> MCOp lore op) -> op -> MCOp lore op
forall a b. (a -> b) -> a -> b
$ Map VName VName -> op -> op
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs op
op

instance (ASTLore lore, Rename op) => Rename (MCOp lore op) where
  rename :: MCOp lore op -> RenameM (MCOp lore op)
rename (ParOp Maybe (SegOp () lore)
par_op SegOp () lore
op) = Maybe (SegOp () lore) -> SegOp () lore -> MCOp lore op
forall lore op.
Maybe (SegOp () lore) -> SegOp () lore -> MCOp lore op
ParOp (Maybe (SegOp () lore) -> SegOp () lore -> MCOp lore op)
-> RenameM (Maybe (SegOp () lore))
-> RenameM (SegOp () lore -> MCOp lore op)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (SegOp () lore) -> RenameM (Maybe (SegOp () lore))
forall a. Rename a => a -> RenameM a
rename Maybe (SegOp () lore)
par_op RenameM (SegOp () lore -> MCOp lore op)
-> RenameM (SegOp () lore) -> RenameM (MCOp lore op)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SegOp () lore -> RenameM (SegOp () lore)
forall a. Rename a => a -> RenameM a
rename SegOp () lore
op
  rename (OtherOp op
op) = op -> MCOp lore op
forall lore op. op -> MCOp lore op
OtherOp (op -> MCOp lore op) -> RenameM op -> RenameM (MCOp lore op)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> op -> RenameM op
forall a. Rename a => a -> RenameM a
rename op
op

instance (ASTLore lore, FreeIn op) => FreeIn (MCOp lore op) where
  freeIn' :: MCOp lore op -> FV
freeIn' (ParOp Maybe (SegOp () lore)
par_op SegOp () lore
op) = Maybe (SegOp () lore) -> FV
forall a. FreeIn a => a -> FV
freeIn' Maybe (SegOp () lore)
par_op FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> SegOp () lore -> FV
forall a. FreeIn a => a -> FV
freeIn' SegOp () lore
op
  freeIn' (OtherOp op
op) = op -> FV
forall a. FreeIn a => a -> FV
freeIn' op
op

instance (ASTLore lore, IsOp op) => IsOp (MCOp lore op) where
  safeOp :: MCOp lore op -> Bool
safeOp (ParOp Maybe (SegOp () lore)
_ SegOp () lore
op) = SegOp () lore -> Bool
forall op. IsOp op => op -> Bool
safeOp SegOp () lore
op
  safeOp (OtherOp op
op) = op -> Bool
forall op. IsOp op => op -> Bool
safeOp op
op

  cheapOp :: MCOp lore op -> Bool
cheapOp (ParOp Maybe (SegOp () lore)
_ SegOp () lore
op) = SegOp () lore -> Bool
forall op. IsOp op => op -> Bool
cheapOp SegOp () lore
op
  cheapOp (OtherOp op
op) = op -> Bool
forall op. IsOp op => op -> Bool
cheapOp op
op

instance TypedOp op => TypedOp (MCOp lore op) where
  opType :: MCOp lore op -> m [ExtType]
opType (ParOp Maybe (SegOp () lore)
_ SegOp () lore
op) = SegOp () lore -> m [ExtType]
forall op t (m :: * -> *).
(TypedOp op, HasScope t m) =>
op -> m [ExtType]
opType SegOp () lore
op
  opType (OtherOp op
op) = op -> m [ExtType]
forall op t (m :: * -> *).
(TypedOp op, HasScope t m) =>
op -> m [ExtType]
opType op
op

instance
  (Aliased lore, AliasedOp op, ASTLore lore) =>
  AliasedOp (MCOp lore op)
  where
  opAliases :: MCOp lore op -> [Names]
opAliases (ParOp Maybe (SegOp () lore)
_ SegOp () lore
op) = SegOp () lore -> [Names]
forall op. AliasedOp op => op -> [Names]
opAliases SegOp () lore
op
  opAliases (OtherOp op
op) = op -> [Names]
forall op. AliasedOp op => op -> [Names]
opAliases op
op

  consumedInOp :: MCOp lore op -> Names
consumedInOp (ParOp Maybe (SegOp () lore)
_ SegOp () lore
op) = SegOp () lore -> Names
forall op. AliasedOp op => op -> Names
consumedInOp SegOp () lore
op
  consumedInOp (OtherOp op
op) = op -> Names
forall op. AliasedOp op => op -> Names
consumedInOp op
op

instance
  (CanBeAliased (Op lore), CanBeAliased op, ASTLore lore) =>
  CanBeAliased (MCOp lore op)
  where
  type OpWithAliases (MCOp lore op) = MCOp (Aliases lore) (OpWithAliases op)

  addOpAliases :: MCOp lore op -> OpWithAliases (MCOp lore op)
addOpAliases (ParOp Maybe (SegOp () lore)
par_op SegOp () lore
op) =
    Maybe (SegOp () (Aliases lore))
-> SegOp () (Aliases lore)
-> MCOp (Aliases lore) (OpWithAliases op)
forall lore op.
Maybe (SegOp () lore) -> SegOp () lore -> MCOp lore op
ParOp (SegOp () lore -> SegOp () (Aliases lore)
forall op. CanBeAliased op => op -> OpWithAliases op
addOpAliases (SegOp () lore -> SegOp () (Aliases lore))
-> Maybe (SegOp () lore) -> Maybe (SegOp () (Aliases lore))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (SegOp () lore)
par_op) (SegOp () lore -> OpWithAliases (SegOp () lore)
forall op. CanBeAliased op => op -> OpWithAliases op
addOpAliases SegOp () lore
op)
  addOpAliases (OtherOp op
op) =
    OpWithAliases op -> MCOp (Aliases lore) (OpWithAliases op)
forall lore op. op -> MCOp lore op
OtherOp (OpWithAliases op -> MCOp (Aliases lore) (OpWithAliases op))
-> OpWithAliases op -> MCOp (Aliases lore) (OpWithAliases op)
forall a b. (a -> b) -> a -> b
$ op -> OpWithAliases op
forall op. CanBeAliased op => op -> OpWithAliases op
addOpAliases op
op

  removeOpAliases :: OpWithAliases (MCOp lore op) -> MCOp lore op
removeOpAliases (ParOp par_op op) =
    Maybe (SegOp () lore) -> SegOp () lore -> MCOp lore op
forall lore op.
Maybe (SegOp () lore) -> SegOp () lore -> MCOp lore op
ParOp (SegOp () (Aliases lore) -> SegOp () lore
forall op. CanBeAliased op => OpWithAliases op -> op
removeOpAliases (SegOp () (Aliases lore) -> SegOp () lore)
-> Maybe (SegOp () (Aliases lore)) -> Maybe (SegOp () lore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (SegOp () (Aliases lore))
par_op) (OpWithAliases (SegOp () lore) -> SegOp () lore
forall op. CanBeAliased op => OpWithAliases op -> op
removeOpAliases OpWithAliases (SegOp () lore)
SegOp () (Aliases lore)
op)
  removeOpAliases (OtherOp op) =
    op -> MCOp lore op
forall lore op. op -> MCOp lore op
OtherOp (op -> MCOp lore op) -> op -> MCOp lore op
forall a b. (a -> b) -> a -> b
$ OpWithAliases op -> op
forall op. CanBeAliased op => OpWithAliases op -> op
removeOpAliases OpWithAliases op
op

instance
  (CanBeWise (Op lore), CanBeWise op, ASTLore lore) =>
  CanBeWise (MCOp lore op)
  where
  type OpWithWisdom (MCOp lore op) = MCOp (Wise lore) (OpWithWisdom op)

  removeOpWisdom :: OpWithWisdom (MCOp lore op) -> MCOp lore op
removeOpWisdom (ParOp par_op op) =
    Maybe (SegOp () lore) -> SegOp () lore -> MCOp lore op
forall lore op.
Maybe (SegOp () lore) -> SegOp () lore -> MCOp lore op
ParOp (SegOp () (Wise lore) -> SegOp () lore
forall op. CanBeWise op => OpWithWisdom op -> op
removeOpWisdom (SegOp () (Wise lore) -> SegOp () lore)
-> Maybe (SegOp () (Wise lore)) -> Maybe (SegOp () lore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (SegOp () (Wise lore))
par_op) (OpWithWisdom (SegOp () lore) -> SegOp () lore
forall op. CanBeWise op => OpWithWisdom op -> op
removeOpWisdom OpWithWisdom (SegOp () lore)
SegOp () (Wise lore)
op)
  removeOpWisdom (OtherOp op) =
    op -> MCOp lore op
forall lore op. op -> MCOp lore op
OtherOp (op -> MCOp lore op) -> op -> MCOp lore op
forall a b. (a -> b) -> a -> b
$ OpWithWisdom op -> op
forall op. CanBeWise op => OpWithWisdom op -> op
removeOpWisdom OpWithWisdom op
op

instance (ASTLore lore, ST.IndexOp op) => ST.IndexOp (MCOp lore op) where
  indexOp :: SymbolTable lore
-> Int -> MCOp lore op -> [TPrimExp Int64 VName] -> Maybe Indexed
indexOp SymbolTable lore
vtable Int
k (ParOp Maybe (SegOp () lore)
_ SegOp () lore
op) [TPrimExp Int64 VName]
is = SymbolTable lore
-> Int -> SegOp () lore -> [TPrimExp Int64 VName] -> Maybe Indexed
forall op lore.
(IndexOp op, ASTLore lore, IndexOp (Op lore)) =>
SymbolTable lore
-> Int -> op -> [TPrimExp Int64 VName] -> Maybe Indexed
ST.indexOp SymbolTable lore
vtable Int
k SegOp () lore
op [TPrimExp Int64 VName]
is
  indexOp SymbolTable lore
vtable Int
k (OtherOp op
op) [TPrimExp Int64 VName]
is = SymbolTable lore
-> Int -> op -> [TPrimExp Int64 VName] -> Maybe Indexed
forall op lore.
(IndexOp op, ASTLore lore, IndexOp (Op lore)) =>
SymbolTable lore
-> Int -> op -> [TPrimExp Int64 VName] -> Maybe Indexed
ST.indexOp SymbolTable lore
vtable Int
k op
op [TPrimExp Int64 VName]
is

instance (PrettyLore lore, Pretty op) => Pretty (MCOp lore op) where
  ppr :: MCOp lore op -> Doc
ppr (ParOp Maybe (SegOp () lore)
Nothing SegOp () lore
op) = SegOp () lore -> Doc
forall a. Pretty a => a -> Doc
ppr SegOp () lore
op
  ppr (ParOp (Just SegOp () lore
par_op) SegOp () lore
op) =
    Doc
"par" Doc -> Doc -> Doc
<+> String -> String -> Doc -> Doc
nestedBlock String
"{" String
"}" (SegOp () lore -> Doc
forall a. Pretty a => a -> Doc
ppr SegOp () lore
par_op)
      Doc -> Doc -> Doc
</> Doc
"seq" Doc -> Doc -> Doc
<+> String -> String -> Doc -> Doc
nestedBlock String
"{" String
"}" (SegOp () lore -> Doc
forall a. Pretty a => a -> Doc
ppr SegOp () lore
op)
  ppr (OtherOp op
op) = op -> Doc
forall a. Pretty a => a -> Doc
ppr op
op

instance (OpMetrics (Op lore), OpMetrics op) => OpMetrics (MCOp lore op) where
  opMetrics :: MCOp lore op -> MetricsM ()
opMetrics (ParOp Maybe (SegOp () lore)
par_op SegOp () lore
op) = Maybe (SegOp () lore) -> MetricsM ()
forall op. OpMetrics op => op -> MetricsM ()
opMetrics Maybe (SegOp () lore)
par_op MetricsM () -> MetricsM () -> MetricsM ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> SegOp () lore -> MetricsM ()
forall op. OpMetrics op => op -> MetricsM ()
opMetrics SegOp () lore
op
  opMetrics (OtherOp op
op) = op -> MetricsM ()
forall op. OpMetrics op => op -> MetricsM ()
opMetrics op
op

typeCheckMCOp ::
  TC.Checkable lore =>
  (op -> TC.TypeM lore ()) ->
  MCOp (Aliases lore) op ->
  TC.TypeM lore ()
typeCheckMCOp :: (op -> TypeM lore ()) -> MCOp (Aliases lore) op -> TypeM lore ()
typeCheckMCOp op -> TypeM lore ()
_ (ParOp (Just SegOp () (Aliases lore)
par_op) SegOp () (Aliases lore)
op) = do
  -- It is valid for the same array to be consumed in both par_op and op.
  ((), ())
_ <- (() -> TypeM lore ()) -> SegOp () (Aliases lore) -> TypeM lore ()
forall lore lvl.
Checkable lore =>
(lvl -> TypeM lore ()) -> SegOp lvl (Aliases lore) -> TypeM lore ()
typeCheckSegOp () -> TypeM lore ()
forall (m :: * -> *) a. Monad m => a -> m a
return SegOp () (Aliases lore)
par_op TypeM lore () -> TypeM lore () -> TypeM lore ((), ())
forall lore a b. TypeM lore a -> TypeM lore b -> TypeM lore (a, b)
`TC.alternative` (() -> TypeM lore ()) -> SegOp () (Aliases lore) -> TypeM lore ()
forall lore lvl.
Checkable lore =>
(lvl -> TypeM lore ()) -> SegOp lvl (Aliases lore) -> TypeM lore ()
typeCheckSegOp () -> TypeM lore ()
forall (m :: * -> *) a. Monad m => a -> m a
return SegOp () (Aliases lore)
op
  () -> TypeM lore ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
typeCheckMCOp op -> TypeM lore ()
_ (ParOp Maybe (SegOp () (Aliases lore))
Nothing SegOp () (Aliases lore)
op) =
  (() -> TypeM lore ()) -> SegOp () (Aliases lore) -> TypeM lore ()
forall lore lvl.
Checkable lore =>
(lvl -> TypeM lore ()) -> SegOp lvl (Aliases lore) -> TypeM lore ()
typeCheckSegOp () -> TypeM lore ()
forall (m :: * -> *) a. Monad m => a -> m a
return SegOp () (Aliases lore)
op
typeCheckMCOp op -> TypeM lore ()
f (OtherOp op
op) = op -> TypeM lore ()
f op
op

simplifyMCOp ::
  ( Engine.SimplifiableLore lore,
    BodyDec lore ~ ()
  ) =>
  Simplify.SimplifyOp lore op ->
  MCOp lore op ->
  Engine.SimpleM lore (MCOp (Wise lore) (OpWithWisdom op), Stms (Wise lore))
simplifyMCOp :: SimplifyOp lore op
-> MCOp lore op
-> SimpleM
     lore (MCOp (Wise lore) (OpWithWisdom op), Stms (Wise lore))
simplifyMCOp SimplifyOp lore op
f (OtherOp op
op) = do
  (OpWithWisdom op
op', Stms (Wise lore)
stms) <- SimplifyOp lore op
f op
op
  (MCOp (Wise lore) (OpWithWisdom op), Stms (Wise lore))
-> SimpleM
     lore (MCOp (Wise lore) (OpWithWisdom op), Stms (Wise lore))
forall (m :: * -> *) a. Monad m => a -> m a
return (OpWithWisdom op -> MCOp (Wise lore) (OpWithWisdom op)
forall lore op. op -> MCOp lore op
OtherOp OpWithWisdom op
op', Stms (Wise lore)
stms)
simplifyMCOp SimplifyOp lore op
_ (ParOp Maybe (SegOp () lore)
par_op SegOp () lore
op) = do
  (Maybe (SegOp () (Wise lore))
par_op', Stms (Wise lore)
par_op_hoisted) <-
    case Maybe (SegOp () lore)
par_op of
      Maybe (SegOp () lore)
Nothing -> (Maybe (SegOp () (Wise lore)), Stms (Wise lore))
-> SimpleM lore (Maybe (SegOp () (Wise lore)), Stms (Wise lore))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (SegOp () (Wise lore))
forall a. Maybe a
Nothing, Stms (Wise lore)
forall a. Monoid a => a
mempty)
      Just SegOp () lore
x -> (SegOp () (Wise lore) -> Maybe (SegOp () (Wise lore)))
-> (SegOp () (Wise lore), Stms (Wise lore))
-> (Maybe (SegOp () (Wise lore)), Stms (Wise lore))
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first SegOp () (Wise lore) -> Maybe (SegOp () (Wise lore))
forall a. a -> Maybe a
Just ((SegOp () (Wise lore), Stms (Wise lore))
 -> (Maybe (SegOp () (Wise lore)), Stms (Wise lore)))
-> SimpleM lore (SegOp () (Wise lore), Stms (Wise lore))
-> SimpleM lore (Maybe (SegOp () (Wise lore)), Stms (Wise lore))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegOp () lore
-> SimpleM lore (SegOp () (Wise lore), Stms (Wise lore))
forall lore lvl.
(SimplifiableLore lore, BodyDec lore ~ (), Simplifiable lvl) =>
SegOp lvl lore
-> SimpleM lore (SegOp lvl (Wise lore), Stms (Wise lore))
simplifySegOp SegOp () lore
x

  (SegOp () (Wise lore)
op', Stms (Wise lore)
op_hoisted) <- SegOp () lore
-> SimpleM lore (SegOp () (Wise lore), Stms (Wise lore))
forall lore lvl.
(SimplifiableLore lore, BodyDec lore ~ (), Simplifiable lvl) =>
SegOp lvl lore
-> SimpleM lore (SegOp lvl (Wise lore), Stms (Wise lore))
simplifySegOp SegOp () lore
op

  (MCOp (Wise lore) (OpWithWisdom op), Stms (Wise lore))
-> SimpleM
     lore (MCOp (Wise lore) (OpWithWisdom op), Stms (Wise lore))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (SegOp () (Wise lore))
-> SegOp () (Wise lore) -> MCOp (Wise lore) (OpWithWisdom op)
forall lore op.
Maybe (SegOp () lore) -> SegOp () lore -> MCOp lore op
ParOp Maybe (SegOp () (Wise lore))
par_op' SegOp () (Wise lore)
op', Stms (Wise lore)
par_op_hoisted Stms (Wise lore) -> Stms (Wise lore) -> Stms (Wise lore)
forall a. Semigroup a => a -> a -> a
<> Stms (Wise lore)
op_hoisted)