-- | Compute a mapping from variables to their corresponding (fully
-- expanded) PrimExps.
module Futhark.Analysis.PrimExp.Table
  ( primExpTable,
    PrimExpTable,

    -- * Extensibility
    PrimExpAnalysis (..),

    -- * Testing
    stmToPrimExps,
  )
where

import Control.Monad.State.Strict
import Data.Foldable
import Data.Map.Strict qualified as M
import Futhark.Analysis.PrimExp
import Futhark.Analysis.PrimExp.Convert
import Futhark.IR.Aliases
import Futhark.IR.GPU
import Futhark.IR.GPUMem
import Futhark.IR.MC
import Futhark.IR.MCMem

-- | Maps variables to maybe PrimExps. Will map to nothing if it
-- cannot be resolved to a PrimExp. For all uses of this analysis atm.
-- a variable can be considered inscrutable if it cannot be resolved
-- to a primexp.
type PrimExpTable = M.Map VName (Maybe (PrimExp VName))

-- | A class for extracting PrimExps from what is inside an op.
class PrimExpAnalysis rep where
  opPrimExp :: Scope rep -> Op rep -> State PrimExpTable ()

primExpTable :: (PrimExpAnalysis rep, RepTypes rep) => Prog rep -> PrimExpTable
primExpTable :: forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Prog rep -> PrimExpTable
primExpTable Prog rep
prog = PrimExpTable
initialState PrimExpTable -> PrimExpTable -> PrimExpTable
forall a. Semigroup a => a -> a -> a
<> ((Scope rep, FunDef rep) -> PrimExpTable)
-> [(Scope rep, FunDef rep)] -> PrimExpTable
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap' ((Scope rep -> FunDef rep -> PrimExpTable)
-> (Scope rep, FunDef rep) -> PrimExpTable
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Scope rep -> FunDef rep -> PrimExpTable
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> FunDef rep -> PrimExpTable
funToPrimExp) [(Scope rep, FunDef rep)]
scopesAndFuns
  where
    scopesAndFuns :: [(Scope rep, FunDef rep)]
scopesAndFuns = do
      let fun_defs :: [FunDef rep]
fun_defs = Prog rep -> [FunDef rep]
forall rep. Prog rep -> [FunDef rep]
progFuns Prog rep
prog
      let scopes :: [Scope rep]
scopes = (FunDef rep -> Scope rep) -> [FunDef rep] -> [Scope rep]
forall a b. (a -> b) -> [a] -> [b]
map FunDef rep -> Scope rep
forall {rep}.
(FParamInfo rep ~ FParamInfo rep) =>
FunDef rep -> Scope rep
getScope [FunDef rep]
fun_defs
      [Scope rep] -> [FunDef rep] -> [(Scope rep, FunDef rep)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Scope rep]
scopes [FunDef rep]
fun_defs

    getScope :: FunDef rep -> Scope rep
getScope FunDef rep
funDef = Stms rep -> Scope rep
forall rep a. Scoped rep a => a -> Scope rep
scopeOf (Prog rep -> Stms rep
forall rep. Prog rep -> Stms rep
progConsts Prog rep
prog) Scope rep -> Scope rep -> Scope rep
forall a. Semigroup a => a -> a -> a
<> [Param (FParamInfo rep)] -> Scope rep
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams (FunDef rep -> [FParam rep]
forall rep. FunDef rep -> [FParam rep]
funDefParams FunDef rep
funDef)

    -- We need to have the dummy "slice" in the analysis for our "slice hack".
    initialState :: PrimExpTable
initialState =
      VName -> Maybe (PrimExp VName) -> PrimExpTable
forall k a. k -> a -> Map k a
M.singleton (Name -> Int -> VName
VName Name
"slice" Int
0) (Maybe (PrimExp VName) -> PrimExpTable)
-> Maybe (PrimExp VName) -> PrimExpTable
forall a b. (a -> b) -> a -> b
$ PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (PrimExp VName -> Maybe (PrimExp VName))
-> PrimExp VName -> Maybe (PrimExp VName)
forall a b. (a -> b) -> a -> b
$ VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp (Name -> Int -> VName
VName Name
"slice" Int
0) (PrimType -> PrimExp VName) -> PrimType -> PrimExp VName
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
Int64

funToPrimExp ::
  (PrimExpAnalysis rep, RepTypes rep) =>
  Scope rep ->
  FunDef rep ->
  PrimExpTable
funToPrimExp :: forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> FunDef rep -> PrimExpTable
funToPrimExp Scope rep
scope FunDef rep
fundef = State PrimExpTable () -> PrimExpTable -> PrimExpTable
forall s a. State s a -> s -> s
execState (Scope rep -> Body rep -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Body rep -> State PrimExpTable ()
bodyToPrimExps Scope rep
scope (FunDef rep -> Body rep
forall rep. FunDef rep -> Body rep
funDefBody FunDef rep
fundef)) PrimExpTable
forall a. Monoid a => a
mempty

-- | Adds the statements of a body to the PrimExpTable
bodyToPrimExps ::
  (PrimExpAnalysis rep, RepTypes rep) =>
  Scope rep ->
  Body rep ->
  State PrimExpTable ()
bodyToPrimExps :: forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Body rep -> State PrimExpTable ()
bodyToPrimExps Scope rep
scope Body rep
body = (Stm rep -> State PrimExpTable ())
-> Seq (Stm rep) -> State PrimExpTable ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Scope rep -> Stm rep -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Stm rep -> State PrimExpTable ()
stmToPrimExps Scope rep
scope') (Body rep -> Seq (Stm rep)
forall rep. Body rep -> Stms rep
bodyStms Body rep
body)
  where
    scope' :: Scope rep
scope' = Scope rep
scope Scope rep -> Scope rep -> Scope rep
forall a. Semigroup a => a -> a -> a
<> Seq (Stm rep) -> Scope rep
forall rep a. Scoped rep a => a -> Scope rep
scopeOf (Body rep -> Seq (Stm rep)
forall rep. Body rep -> Stms rep
bodyStms Body rep
body)

-- | Adds the statements of a kernel body to the PrimExpTable
kernelToBodyPrimExps ::
  (PrimExpAnalysis rep, RepTypes rep) =>
  Scope rep ->
  KernelBody rep ->
  State PrimExpTable ()
kernelToBodyPrimExps :: forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> KernelBody rep -> State PrimExpTable ()
kernelToBodyPrimExps Scope rep
scope KernelBody rep
kbody = (Stm rep -> State PrimExpTable ())
-> Seq (Stm rep) -> State PrimExpTable ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Scope rep -> Stm rep -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Stm rep -> State PrimExpTable ()
stmToPrimExps Scope rep
scope') (KernelBody rep -> Seq (Stm rep)
forall rep. KernelBody rep -> Stms rep
kernelBodyStms KernelBody rep
kbody)
  where
    scope' :: Scope rep
scope' = Scope rep
scope Scope rep -> Scope rep -> Scope rep
forall a. Semigroup a => a -> a -> a
<> Seq (Stm rep) -> Scope rep
forall rep a. Scoped rep a => a -> Scope rep
scopeOf (KernelBody rep -> Seq (Stm rep)
forall rep. KernelBody rep -> Stms rep
kernelBodyStms KernelBody rep
kbody)

-- | Adds a statement to the PrimExpTable. If it can't be resolved as a `PrimExp`,
-- it will be added as `Nothing`.
stmToPrimExps ::
  forall rep.
  (PrimExpAnalysis rep, RepTypes rep) =>
  Scope rep ->
  Stm rep ->
  State PrimExpTable ()
stmToPrimExps :: forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Stm rep -> State PrimExpTable ()
stmToPrimExps Scope rep
scope Stm rep
stm = do
  PrimExpTable
table <- StateT PrimExpTable Identity PrimExpTable
forall s (m :: * -> *). MonadState s m => m s
get
  case Stm rep
stm of
    (Let (Pat [PatElem (LetDec rep)]
pat_elems) StmAux (ExpDec rep)
_ Exp rep
e)
      | Just PrimExp VName
primExp <- (VName -> Maybe (PrimExp VName))
-> Exp rep -> Maybe (PrimExp VName)
forall (m :: * -> *) rep v.
(MonadFail m, RepTypes rep) =>
(VName -> m (PrimExp v)) -> Exp rep -> m (PrimExp v)
primExpFromExp (Scope rep -> PrimExpTable -> VName -> Maybe (PrimExp VName)
forall rep.
RepTypes rep =>
Scope rep -> PrimExpTable -> VName -> Maybe (PrimExp VName)
toPrimExp Scope rep
scope PrimExpTable
table) Exp rep
e ->
          -- The statement can be resolved as a `PrimExp`.
          -- For each pattern element, insert the PrimExp in the table
          [PatElem (LetDec rep)]
-> (PatElem (LetDec rep) -> State PrimExpTable ())
-> State PrimExpTable ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [PatElem (LetDec rep)]
pat_elems ((PatElem (LetDec rep) -> State PrimExpTable ())
 -> State PrimExpTable ())
-> (PatElem (LetDec rep) -> State PrimExpTable ())
-> State PrimExpTable ()
forall a b. (a -> b) -> a -> b
$ \PatElem (LetDec rep)
pe ->
            (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((PrimExpTable -> PrimExpTable) -> State PrimExpTable ())
-> (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
forall a b. (a -> b) -> a -> b
$ VName -> Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (LetDec rep)
pe) (PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just PrimExp VName
primExp)
      | Bool
otherwise -> do
          -- The statement can't be resolved as a `PrimExp`.
          Exp rep -> State PrimExpTable ()
walk (Exp rep -> State PrimExpTable ())
-> Exp rep -> State PrimExpTable ()
forall a b. (a -> b) -> a -> b
$ Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm -- Traverse the rest of the AST Get the
          -- updated PrimExpTable after traversing the AST
          PrimExpTable
table' <- StateT PrimExpTable Identity PrimExpTable
forall s (m :: * -> *). MonadState s m => m s
get

          -- Add pattern elements that can't be resolved as `PrimExp`
          -- to the `PrimExpTable` as `Nothing`
          [PatElem (LetDec rep)]
-> (PatElem (LetDec rep) -> State PrimExpTable ())
-> State PrimExpTable ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [PatElem (LetDec rep)]
pat_elems ((PatElem (LetDec rep) -> State PrimExpTable ())
 -> State PrimExpTable ())
-> (PatElem (LetDec rep) -> State PrimExpTable ())
-> State PrimExpTable ()
forall a b. (a -> b) -> a -> b
$ \PatElem (LetDec rep)
pe ->
            case VName -> PrimExpTable -> Maybe (Maybe (PrimExp VName))
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (LetDec rep)
pe) PrimExpTable
table' of
              Maybe (Maybe (PrimExp VName))
Nothing -> (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((PrimExpTable -> PrimExpTable) -> State PrimExpTable ())
-> (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
forall a b. (a -> b) -> a -> b
$ VName -> Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (LetDec rep)
pe) Maybe (PrimExp VName)
forall a. Maybe a
Nothing
              Just Maybe (PrimExp VName)
_ -> () -> State PrimExpTable ()
forall a. a -> StateT PrimExpTable Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  where
    walk :: Exp rep -> State PrimExpTable ()
walk Exp rep
e = do
      -- Handle most cases using the walker
      Walker rep (StateT PrimExpTable Identity)
-> Exp rep -> State PrimExpTable ()
forall (m :: * -> *) rep.
Monad m =>
Walker rep m -> Exp rep -> m ()
walkExpM Walker rep (StateT PrimExpTable Identity)
walker Exp rep
e
      -- Additionally, handle loop parameters
      case Exp rep
e of
        Loop [(FParam rep, SubExp)]
_ (ForLoop VName
i IntType
t SubExp
_) Body rep
_ ->
          (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((PrimExpTable -> PrimExpTable) -> State PrimExpTable ())
-> (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
forall a b. (a -> b) -> a -> b
$ VName -> Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
i (Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable)
-> Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable
forall a b. (a -> b) -> a -> b
$ PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (PrimExp VName -> Maybe (PrimExp VName))
-> PrimExp VName -> Maybe (PrimExp VName)
forall a b. (a -> b) -> a -> b
$ VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
i (PrimType -> PrimExp VName) -> PrimType -> PrimExp VName
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
t
        Exp rep
_ -> () -> State PrimExpTable ()
forall a. a -> StateT PrimExpTable Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

    walker :: Walker rep (StateT PrimExpTable Identity)
walker =
      (forall rep (m :: * -> *). Monad m => Walker rep m
identityWalker @rep)
        { walkOnBody = \Scope rep
body_scope -> Scope rep -> Body rep -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Body rep -> State PrimExpTable ()
bodyToPrimExps (Scope rep
scope Scope rep -> Scope rep -> Scope rep
forall a. Semigroup a => a -> a -> a
<> Scope rep
body_scope),
          walkOnOp = opPrimExp scope,
          walkOnFParam = paramToPrimExp -- Loop parameters
        }

    -- Adds a loop parameter to the PrimExpTable
    paramToPrimExp :: FParam rep -> State PrimExpTable ()
    paramToPrimExp :: FParam rep -> State PrimExpTable ()
paramToPrimExp FParam rep
param = do
      let name :: VName
name = FParam rep -> VName
forall dec. Param dec -> VName
paramName FParam rep
param
      -- Construct a `PrimExp` from the type of the parameter
      -- and add it to the `PrimExpTable`
      case FParamInfo rep -> Type
forall t. Typed t => t -> Type
typeOf (FParamInfo rep -> Type) -> FParamInfo rep -> Type
forall a b. (a -> b) -> a -> b
$ FParam rep -> FParamInfo rep
forall dec. Param dec -> dec
paramDec FParam rep
param of
        -- TODO: Handle other types?
        Prim PrimType
pt ->
          (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((PrimExpTable -> PrimExpTable) -> State PrimExpTable ())
-> (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
forall a b. (a -> b) -> a -> b
$ VName -> Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
name (PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (PrimExp VName -> Maybe (PrimExp VName))
-> PrimExp VName -> Maybe (PrimExp VName)
forall a b. (a -> b) -> a -> b
$ VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
name PrimType
pt)
        Type
_ -> () -> State PrimExpTable ()
forall a. a -> StateT PrimExpTable Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Checks if a name is in the PrimExpTable and construct a `PrimExp`
-- if it is not
toPrimExp :: (RepTypes rep) => Scope rep -> PrimExpTable -> VName -> Maybe (PrimExp VName)
toPrimExp :: forall rep.
RepTypes rep =>
Scope rep -> PrimExpTable -> VName -> Maybe (PrimExp VName)
toPrimExp Scope rep
scope PrimExpTable
table VName
name = case VName -> PrimExpTable -> Maybe (Maybe (PrimExp VName))
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name PrimExpTable
table of
  Just Maybe (PrimExp VName)
maybePrimExp
    | Just PrimExp VName
primExp <- Maybe (PrimExp VName)
maybePrimExp -> PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just PrimExp VName
primExp -- Already in the table
  Maybe (Maybe (PrimExp VName))
_ -> case (NameInfo rep -> Type) -> Maybe (NameInfo rep) -> Maybe Type
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NameInfo rep -> Type
forall t. Typed t => t -> Type
typeOf (Maybe (NameInfo rep) -> Maybe Type)
-> (Scope rep -> Maybe (NameInfo rep)) -> Scope rep -> Maybe Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Scope rep -> Maybe (NameInfo rep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (Scope rep -> Maybe Type) -> Scope rep -> Maybe Type
forall a b. (a -> b) -> a -> b
$ Scope rep
scope of
    (Just (Prim PrimType
pt)) -> PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (PrimExp VName -> Maybe (PrimExp VName))
-> PrimExp VName -> Maybe (PrimExp VName)
forall a b. (a -> b) -> a -> b
$ VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
name PrimType
pt
    Maybe Type
_ -> Maybe (PrimExp VName)
forall a. Maybe a
Nothing

-- | Adds the parameters of a SegOp as well as the statements in its
-- body to the PrimExpTable
segOpToPrimExps :: (PrimExpAnalysis rep, RepTypes rep) => Scope rep -> SegOp lvl rep -> State PrimExpTable ()
segOpToPrimExps :: forall rep lvl.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> SegOp lvl rep -> State PrimExpTable ()
segOpToPrimExps Scope rep
scope SegOp lvl rep
op = do
  [VName]
-> (VName -> State PrimExpTable ()) -> State PrimExpTable ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (((VName, SubExp) -> VName) -> [(VName, SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> VName
forall a b. (a, b) -> a
fst ([(VName, SubExp)] -> [VName]) -> [(VName, SubExp)] -> [VName]
forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace (SegSpace -> [(VName, SubExp)]) -> SegSpace -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$ SegOp lvl rep -> SegSpace
forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp lvl rep
op) ((VName -> State PrimExpTable ()) -> State PrimExpTable ())
-> (VName -> State PrimExpTable ()) -> State PrimExpTable ()
forall a b. (a -> b) -> a -> b
$ \VName
name ->
    (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((PrimExpTable -> PrimExpTable) -> State PrimExpTable ())
-> (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
forall a b. (a -> b) -> a -> b
$ VName -> Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
name (Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable)
-> Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable
forall a b. (a -> b) -> a -> b
$ PrimExp VName -> Maybe (PrimExp VName)
forall a. a -> Maybe a
Just (PrimExp VName -> Maybe (PrimExp VName))
-> PrimExp VName -> Maybe (PrimExp VName)
forall a b. (a -> b) -> a -> b
$ VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
name PrimType
int64
  Scope rep -> KernelBody rep -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> KernelBody rep -> State PrimExpTable ()
kernelToBodyPrimExps Scope rep
scope (SegOp lvl rep -> KernelBody rep
forall lvl rep. SegOp lvl rep -> KernelBody rep
segBody SegOp lvl rep
op)

instance PrimExpAnalysis GPU where
  opPrimExp :: Scope GPU -> Op GPU -> State PrimExpTable ()
opPrimExp Scope GPU
scope Op GPU
gpu_op
    | (SegOp SegOp SegLevel GPU
op) <- Op GPU
gpu_op = Scope GPU -> SegOp SegLevel GPU -> State PrimExpTable ()
forall rep lvl.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> SegOp lvl rep -> State PrimExpTable ()
segOpToPrimExps Scope GPU
scope SegOp SegLevel GPU
op
    | (SizeOp SizeOp
_) <- Op GPU
gpu_op = () -> State PrimExpTable ()
forall a. a -> StateT PrimExpTable Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    | (GPUBody [Type]
_ Body GPU
body) <- Op GPU
gpu_op = Scope GPU -> Body GPU -> State PrimExpTable ()
forall rep.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> Body rep -> State PrimExpTable ()
bodyToPrimExps Scope GPU
scope Body GPU
body
    | (Futhark.IR.GPUMem.OtherOp SOAC GPU
_) <- Op GPU
gpu_op = () -> State PrimExpTable ()
forall a. a -> StateT PrimExpTable Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

instance PrimExpAnalysis MC where
  opPrimExp :: Scope MC -> Op MC -> State PrimExpTable ()
opPrimExp Scope MC
scope Op MC
mc_op
    | (ParOp Maybe (SegOp () MC)
maybe_par_segop SegOp () MC
seq_segop) <- Op MC
mc_op = do
        -- Add the statements in the parallel part of the ParOp to the PrimExpTable
        case Maybe (SegOp () MC)
maybe_par_segop of
          Maybe (SegOp () MC)
Nothing -> () -> State PrimExpTable ()
forall a. a -> StateT PrimExpTable Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          Just SegOp () MC
_ -> Maybe (SegOp () MC)
-> (SegOp () MC -> State PrimExpTable ()) -> State PrimExpTable ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe (SegOp () MC)
maybe_par_segop ((SegOp () MC -> State PrimExpTable ()) -> State PrimExpTable ())
-> (SegOp () MC -> State PrimExpTable ()) -> State PrimExpTable ()
forall a b. (a -> b) -> a -> b
$ Scope MC -> SegOp () MC -> State PrimExpTable ()
forall rep lvl.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> SegOp lvl rep -> State PrimExpTable ()
segOpToPrimExps Scope MC
scope
        -- Add the statements in the sequential part of the ParOp to the PrimExpTable
        Scope MC -> SegOp () MC -> State PrimExpTable ()
forall rep lvl.
(PrimExpAnalysis rep, RepTypes rep) =>
Scope rep -> SegOp lvl rep -> State PrimExpTable ()
segOpToPrimExps Scope MC
scope SegOp () MC
seq_segop
    | (Futhark.IR.MCMem.OtherOp SOAC MC
_) <- Op MC
mc_op = () -> State PrimExpTable ()
forall a. a -> StateT PrimExpTable Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()