{-# LANGUAGE TypeFamilies #-}

module Futhark.IR.SeqMem
  ( SeqMem,

    -- * Simplification
    simplifyProg,
    simpleSeqMem,

    -- * Module re-exports
    module Futhark.IR.Mem,
  )
where

import Futhark.Analysis.PrimExp.Convert
import Futhark.IR.Mem
import Futhark.IR.Mem.Simplify
import Futhark.IR.TypeCheck qualified as TC
import Futhark.Optimise.Simplify.Engine qualified as Engine
import Futhark.Pass
import Futhark.Pass.ExplicitAllocations (BuilderOps (..), mkLetNamesB', mkLetNamesB'')

data SeqMem

instance RepTypes SeqMem where
  type LetDec SeqMem = LetDecMem
  type FParamInfo SeqMem = FParamMem
  type LParamInfo SeqMem = LParamMem
  type RetType SeqMem = RetTypeMem
  type BranchType SeqMem = BranchTypeMem
  type OpC SeqMem = MemOp NoOp

instance ASTRep SeqMem where
  expTypesFromPat :: forall (m :: * -> *).
(HasScope SeqMem m, Monad m) =>
Pat (LetDec SeqMem) -> m [BranchType SeqMem]
expTypesFromPat = [BranchTypeMem] -> m [BranchTypeMem]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([BranchTypeMem] -> m [BranchTypeMem])
-> (Pat LetDecMem -> [BranchTypeMem])
-> Pat LetDecMem
-> m [BranchTypeMem]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, BranchTypeMem) -> BranchTypeMem)
-> [(VName, BranchTypeMem)] -> [BranchTypeMem]
forall a b. (a -> b) -> [a] -> [b]
map (VName, BranchTypeMem) -> BranchTypeMem
forall a b. (a, b) -> b
snd ([(VName, BranchTypeMem)] -> [BranchTypeMem])
-> (Pat LetDecMem -> [(VName, BranchTypeMem)])
-> Pat LetDecMem
-> [BranchTypeMem]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat LetDecMem -> [(VName, BranchTypeMem)]
bodyReturnsFromPat

instance PrettyRep SeqMem

instance TC.Checkable SeqMem where
  checkOp :: Op (Aliases SeqMem) -> TypeM SeqMem ()
checkOp (Alloc SubExp
size Space
_) = [TypeBase (ShapeBase SubExp) NoUniqueness]
-> SubExp -> TypeM SeqMem ()
forall rep.
Checkable rep =>
[TypeBase (ShapeBase SubExp) NoUniqueness]
-> SubExp -> TypeM rep ()
TC.require [PrimType -> TypeBase (ShapeBase SubExp) NoUniqueness
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int64] SubExp
size
  checkOp (Inner NoOp (Aliases SeqMem)
NoOp) = () -> TypeM SeqMem ()
forall a. a -> TypeM SeqMem a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  checkFParamDec :: VName -> FParamInfo SeqMem -> TypeM SeqMem ()
checkFParamDec = VName -> FParamInfo SeqMem -> TypeM SeqMem ()
VName -> FParamMem -> TypeM SeqMem ()
forall rep u.
Checkable rep =>
VName -> MemInfo SubExp u MemBind -> TypeM rep ()
checkMemInfo
  checkLParamDec :: VName -> LParamInfo SeqMem -> TypeM SeqMem ()
checkLParamDec = VName -> LParamInfo SeqMem -> TypeM SeqMem ()
VName -> LetDecMem -> TypeM SeqMem ()
forall rep u.
Checkable rep =>
VName -> MemInfo SubExp u MemBind -> TypeM rep ()
checkMemInfo
  checkLetBoundDec :: VName -> LetDec SeqMem -> TypeM SeqMem ()
checkLetBoundDec = VName -> LetDec SeqMem -> TypeM SeqMem ()
VName -> LetDecMem -> TypeM SeqMem ()
forall rep u.
Checkable rep =>
VName -> MemInfo SubExp u MemBind -> TypeM rep ()
checkMemInfo
  checkRetType :: [RetType SeqMem] -> TypeM SeqMem ()
checkRetType = (RetTypeMem -> TypeM SeqMem ()) -> [RetTypeMem] -> TypeM SeqMem ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (TypeBase (ShapeBase (Ext SubExp)) Uniqueness -> TypeM SeqMem ()
forall rep u.
Checkable rep =>
TypeBase (ShapeBase (Ext SubExp)) u -> TypeM rep ()
TC.checkExtType (TypeBase (ShapeBase (Ext SubExp)) Uniqueness -> TypeM SeqMem ())
-> (RetTypeMem -> TypeBase (ShapeBase (Ext SubExp)) Uniqueness)
-> RetTypeMem
-> TypeM SeqMem ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RetTypeMem -> TypeBase (ShapeBase (Ext SubExp)) Uniqueness
forall t.
DeclExtTyped t =>
t -> TypeBase (ShapeBase (Ext SubExp)) Uniqueness
declExtTypeOf)
  primFParam :: VName -> PrimType -> TypeM SeqMem (FParam (Aliases SeqMem))
primFParam VName
name PrimType
t = FParam (Aliases SeqMem) -> TypeM SeqMem (FParam (Aliases SeqMem))
forall a. a -> TypeM SeqMem a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FParam (Aliases SeqMem) -> TypeM SeqMem (FParam (Aliases SeqMem)))
-> FParam (Aliases SeqMem)
-> TypeM SeqMem (FParam (Aliases SeqMem))
forall a b. (a -> b) -> a -> b
$ Attrs -> VName -> FParamMem -> Param FParamMem
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
forall a. Monoid a => a
mempty VName
name (PrimType -> FParamMem
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
t)
  matchPat :: Pat (LetDec (Aliases SeqMem))
-> Exp (Aliases SeqMem) -> TypeM SeqMem ()
matchPat = Pat (LetDec (Aliases SeqMem))
-> Exp (Aliases SeqMem) -> TypeM SeqMem ()
forall rep (inner :: * -> *).
(Mem rep inner, LetDec rep ~ LetDecMem, Checkable rep) =>
Pat (LetDec (Aliases rep)) -> Exp (Aliases rep) -> TypeM rep ()
matchPatToExp
  matchReturnType :: [RetType SeqMem] -> Result -> TypeM SeqMem ()
matchReturnType = [RetType SeqMem] -> Result -> TypeM SeqMem ()
[RetTypeMem] -> Result -> TypeM SeqMem ()
forall rep (inner :: * -> *).
(Mem rep inner, Checkable rep) =>
[RetTypeMem] -> Result -> TypeM rep ()
matchFunctionReturnType
  matchBranchType :: [BranchType SeqMem] -> Body (Aliases SeqMem) -> TypeM SeqMem ()
matchBranchType = [BranchType SeqMem] -> Body (Aliases SeqMem) -> TypeM SeqMem ()
[BranchTypeMem] -> Body (Aliases SeqMem) -> TypeM SeqMem ()
forall rep (inner :: * -> *).
(Mem rep inner, Checkable rep) =>
[BranchTypeMem] -> Body (Aliases rep) -> TypeM rep ()
matchBranchReturnType
  matchLoopResult :: [FParam (Aliases SeqMem)] -> Result -> TypeM SeqMem ()
matchLoopResult = [FParam (Aliases SeqMem)] -> Result -> TypeM SeqMem ()
forall rep (inner :: * -> *).
(Mem rep inner, Checkable rep) =>
[FParam (Aliases rep)] -> Result -> TypeM rep ()
matchLoopResultMem

instance BuilderOps SeqMem where
  mkExpDecB :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ SeqMem) =>
Pat (LetDec SeqMem) -> Exp SeqMem -> m (ExpDec SeqMem)
mkExpDecB Pat (LetDec SeqMem)
_ Exp SeqMem
_ = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  mkBodyB :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ SeqMem) =>
Stms SeqMem -> Result -> m (Body SeqMem)
mkBodyB Stms SeqMem
stms Result
res = Body SeqMem -> m (Body SeqMem)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Body SeqMem -> m (Body SeqMem)) -> Body SeqMem -> m (Body SeqMem)
forall a b. (a -> b) -> a -> b
$ BodyDec SeqMem -> Stms SeqMem -> Result -> Body SeqMem
forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body () Stms SeqMem
stms Result
res
  mkLetNamesB :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ SeqMem) =>
[VName] -> Exp SeqMem -> m (Stm SeqMem)
mkLetNamesB = Space
-> ExpDec (Rep m) -> [VName] -> Exp (Rep m) -> m (Stm (Rep m))
forall (m :: * -> *) (inner :: * -> *).
(LetDec (Rep m) ~ LetDecMem, Mem (Rep m) inner, MonadBuilder m,
 ExpDec (Rep m) ~ ()) =>
Space
-> ExpDec (Rep m) -> [VName] -> Exp (Rep m) -> m (Stm (Rep m))
mkLetNamesB' Space
DefaultSpace ()

instance TraverseOpStms SeqMem where
  traverseOpStms :: forall (m :: * -> *).
Monad m =>
OpStmsTraverser m (Op SeqMem) SeqMem
traverseOpStms Scope SeqMem -> Stms SeqMem -> m (Stms SeqMem)
_ = Op SeqMem -> m (Op SeqMem)
MemOp NoOp SeqMem -> m (MemOp NoOp SeqMem)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

instance BuilderOps (Engine.Wise SeqMem) where
  mkExpDecB :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ Wise SeqMem) =>
Pat (LetDec (Wise SeqMem))
-> Exp (Wise SeqMem) -> m (ExpDec (Wise SeqMem))
mkExpDecB Pat (LetDec (Wise SeqMem))
pat Exp (Wise SeqMem)
e = ExpDec (Wise SeqMem) -> m (ExpDec (Wise SeqMem))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpDec (Wise SeqMem) -> m (ExpDec (Wise SeqMem)))
-> ExpDec (Wise SeqMem) -> m (ExpDec (Wise SeqMem))
forall a b. (a -> b) -> a -> b
$ Pat (LetDec (Wise SeqMem))
-> ExpDec SeqMem -> Exp (Wise SeqMem) -> ExpDec (Wise SeqMem)
forall rep.
Informing rep =>
Pat (LetDec (Wise rep))
-> ExpDec rep -> Exp (Wise rep) -> ExpDec (Wise rep)
Engine.mkWiseExpDec Pat (LetDec (Wise SeqMem))
pat () Exp (Wise SeqMem)
e
  mkBodyB :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ Wise SeqMem) =>
Stms (Wise SeqMem) -> Result -> m (Body (Wise SeqMem))
mkBodyB Stms (Wise SeqMem)
stms Result
res = Body (Wise SeqMem) -> m (Body (Wise SeqMem))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Body (Wise SeqMem) -> m (Body (Wise SeqMem)))
-> Body (Wise SeqMem) -> m (Body (Wise SeqMem))
forall a b. (a -> b) -> a -> b
$ BodyDec SeqMem
-> Stms (Wise SeqMem) -> Result -> Body (Wise SeqMem)
forall rep.
Informing rep =>
BodyDec rep -> Stms (Wise rep) -> Result -> Body (Wise rep)
Engine.mkWiseBody () Stms (Wise SeqMem)
stms Result
res
  mkLetNamesB :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ Wise SeqMem) =>
[VName] -> Exp (Wise SeqMem) -> m (Stm (Wise SeqMem))
mkLetNamesB = Space -> [VName] -> Exp (Wise SeqMem) -> m (Stm (Wise SeqMem))
forall rep (inner :: * -> *) (m :: * -> *).
(Mem rep inner, LetDec rep ~ LetDecMem,
 OpReturns (inner (Wise rep)), ExpDec rep ~ (), Rep m ~ Wise rep,
 HasScope (Wise rep) m, MonadBuilder m,
 AliasedOp (inner (Wise rep)), RephraseOp (MemOp inner),
 CanBeWise inner) =>
Space -> [VName] -> Exp (Wise rep) -> m (Stm (Wise rep))
mkLetNamesB'' Space
DefaultSpace

instance TraverseOpStms (Engine.Wise SeqMem) where
  traverseOpStms :: forall (m :: * -> *).
Monad m =>
OpStmsTraverser m (Op (Wise SeqMem)) (Wise SeqMem)
traverseOpStms Scope (Wise SeqMem) -> Stms (Wise SeqMem) -> m (Stms (Wise SeqMem))
_ = Op (Wise SeqMem) -> m (Op (Wise SeqMem))
MemOp NoOp (Wise SeqMem) -> m (MemOp NoOp (Wise SeqMem))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

simplifyProg :: Prog SeqMem -> PassM (Prog SeqMem)
simplifyProg :: Prog SeqMem -> PassM (Prog SeqMem)
simplifyProg = RuleBook (Wise SeqMem)
-> SimpleOps SeqMem -> Prog SeqMem -> PassM (Prog SeqMem)
forall rep (inner :: * -> *).
SimplifyMemory rep inner =>
RuleBook (Wise rep)
-> SimpleOps rep -> Prog rep -> PassM (Prog rep)
simplifyProgGeneric RuleBook (Wise SeqMem)
forall rep (inner :: * -> *).
SimplifyMemory rep inner =>
RuleBook (Wise rep)
memRuleBook SimpleOps SeqMem
simpleSeqMem

simpleSeqMem :: Engine.SimpleOps SeqMem
simpleSeqMem :: SimpleOps SeqMem
simpleSeqMem =
  (NoOp (Wise SeqMem) -> UsageTable)
-> SimplifyOp SeqMem (NoOp (Wise SeqMem)) -> SimpleOps SeqMem
forall rep (inner :: * -> *).
SimplifyMemory rep inner =>
(inner (Wise rep) -> UsageTable)
-> SimplifyOp rep (inner (Wise rep)) -> SimpleOps rep
simpleGeneric (UsageTable -> NoOp (Wise SeqMem) -> UsageTable
forall a b. a -> b -> a
const UsageTable
forall a. Monoid a => a
mempty) (SimplifyOp SeqMem (NoOp (Wise SeqMem)) -> SimpleOps SeqMem)
-> SimplifyOp SeqMem (NoOp (Wise SeqMem)) -> SimpleOps SeqMem
forall a b. (a -> b) -> a -> b
$ SimpleM SeqMem (NoOp (Wise SeqMem), Stms (Wise SeqMem))
-> SimplifyOp SeqMem (NoOp (Wise SeqMem))
forall a b. a -> b -> a
const (SimpleM SeqMem (NoOp (Wise SeqMem), Stms (Wise SeqMem))
 -> SimplifyOp SeqMem (NoOp (Wise SeqMem)))
-> SimpleM SeqMem (NoOp (Wise SeqMem), Stms (Wise SeqMem))
-> SimplifyOp SeqMem (NoOp (Wise SeqMem))
forall a b. (a -> b) -> a -> b
$ (NoOp (Wise SeqMem), Stms (Wise SeqMem))
-> SimpleM SeqMem (NoOp (Wise SeqMem), Stms (Wise SeqMem))
forall a. a -> SimpleM SeqMem a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NoOp (Wise SeqMem)
forall {k} (rep :: k). NoOp rep
NoOp, Stms (Wise SeqMem)
forall a. Monoid a => a
mempty)