{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}

module Futhark.IR.MCMem
  ( MCMem,

    -- * Simplification
    simplifyProg,

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

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

data MCMem

instance RepTypes MCMem where
  type LetDec MCMem = LetDecMem
  type FParamInfo MCMem = FParamMem
  type LParamInfo MCMem = LParamMem
  type RetType MCMem = RetTypeMem
  type BranchType MCMem = BranchTypeMem
  type Op MCMem = MemOp (MCOp MCMem ())

instance ASTRep MCMem where
  expTypesFromPat :: Pat (LetDec MCMem) -> m [BranchType MCMem]
expTypesFromPat = [BodyReturns] -> m [BodyReturns]
forall (m :: * -> *) a. Monad m => a -> m a
return ([BodyReturns] -> m [BodyReturns])
-> (Pat (MemBound NoUniqueness) -> [BodyReturns])
-> Pat (MemBound NoUniqueness)
-> m [BodyReturns]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, BodyReturns) -> BodyReturns)
-> [(VName, BodyReturns)] -> [BodyReturns]
forall a b. (a -> b) -> [a] -> [b]
map (VName, BodyReturns) -> BodyReturns
forall a b. (a, b) -> b
snd ([(VName, BodyReturns)] -> [BodyReturns])
-> (Pat (MemBound NoUniqueness) -> [(VName, BodyReturns)])
-> Pat (MemBound NoUniqueness)
-> [BodyReturns]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (MemBound NoUniqueness) -> [(VName, BodyReturns)]
bodyReturnsFromPat

instance OpReturns (MCOp MCMem ()) where
  opReturns :: MCOp MCMem () -> m [ExpReturns]
opReturns (ParOp Maybe (SegOp () MCMem)
_ SegOp () MCMem
op) = SegOp () MCMem -> m [ExpReturns]
forall rep inner (m :: * -> *) lvl somerep.
(Mem rep inner, Monad m, HasScope rep m) =>
SegOp lvl somerep -> m [ExpReturns]
segOpReturns SegOp () MCMem
op
  opReturns (OtherOp ()) = [ExpReturns] -> m [ExpReturns]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []

instance OpReturns (MCOp (Engine.Wise MCMem) ()) where
  opReturns :: MCOp (Wise MCMem) () -> m [ExpReturns]
opReturns (ParOp Maybe (SegOp () (Wise MCMem))
_ SegOp () (Wise MCMem)
op) = SegOp () (Wise MCMem) -> m [ExpReturns]
forall rep inner (m :: * -> *) lvl somerep.
(Mem rep inner, Monad m, HasScope rep m) =>
SegOp lvl somerep -> m [ExpReturns]
segOpReturns SegOp () (Wise MCMem)
op
  opReturns MCOp (Wise MCMem) ()
k = [ExtType] -> [ExpReturns]
extReturns ([ExtType] -> [ExpReturns]) -> m [ExtType] -> m [ExpReturns]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MCOp (Wise MCMem) () -> m [ExtType]
forall op t (m :: * -> *).
(TypedOp op, HasScope t m) =>
op -> m [ExtType]
opType MCOp (Wise MCMem) ()
k

instance PrettyRep MCMem

instance TC.CheckableOp MCMem where
  checkOp :: OpWithAliases (Op MCMem) -> TypeM MCMem ()
checkOp = OpWithAliases (Op MCMem) -> TypeM MCMem ()
forall rep.
Checkable rep =>
MemOp (MCOp (Aliases rep) ()) -> TypeM rep ()
typeCheckMemoryOp
    where
      typeCheckMemoryOp :: MemOp (MCOp (Aliases rep) ()) -> TypeM rep ()
typeCheckMemoryOp (Alloc SubExp
size Space
_) =
        [Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
TC.require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int64] SubExp
size
      typeCheckMemoryOp (Inner MCOp (Aliases rep) ()
op) =
        (() -> TypeM rep ()) -> MCOp (Aliases rep) () -> TypeM rep ()
forall rep op.
Checkable rep =>
(op -> TypeM rep ()) -> MCOp (Aliases rep) op -> TypeM rep ()
typeCheckMCOp () -> TypeM rep ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure MCOp (Aliases rep) ()
op

instance TC.Checkable MCMem where
  checkFParamDec :: VName -> FParamInfo MCMem -> TypeM MCMem ()
checkFParamDec = VName -> FParamInfo MCMem -> TypeM MCMem ()
forall rep u.
Checkable rep =>
VName -> MemInfo SubExp u MemBind -> TypeM rep ()
checkMemInfo
  checkLParamDec :: VName -> LParamInfo MCMem -> TypeM MCMem ()
checkLParamDec = VName -> LParamInfo MCMem -> TypeM MCMem ()
forall rep u.
Checkable rep =>
VName -> MemInfo SubExp u MemBind -> TypeM rep ()
checkMemInfo
  checkLetBoundDec :: VName -> LetDec MCMem -> TypeM MCMem ()
checkLetBoundDec = VName -> LetDec MCMem -> TypeM MCMem ()
forall rep u.
Checkable rep =>
VName -> MemInfo SubExp u MemBind -> TypeM rep ()
checkMemInfo
  checkRetType :: [RetType MCMem] -> TypeM MCMem ()
checkRetType = (RetTypeMem -> TypeM MCMem ()) -> [RetTypeMem] -> TypeM MCMem ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (TypeBase ExtShape Uniqueness -> TypeM MCMem ()
forall rep u. Checkable rep => TypeBase ExtShape u -> TypeM rep ()
TC.checkExtType (TypeBase ExtShape Uniqueness -> TypeM MCMem ())
-> (RetTypeMem -> TypeBase ExtShape Uniqueness)
-> RetTypeMem
-> TypeM MCMem ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RetTypeMem -> TypeBase ExtShape Uniqueness
forall t. DeclExtTyped t => t -> TypeBase ExtShape Uniqueness
declExtTypeOf)
  primFParam :: VName -> PrimType -> TypeM MCMem (FParam (Aliases MCMem))
primFParam VName
name PrimType
t = Param (MemInfo SubExp Uniqueness MemBind)
-> TypeM MCMem (Param (MemInfo SubExp Uniqueness MemBind))
forall (m :: * -> *) a. Monad m => a -> m a
return (Param (MemInfo SubExp Uniqueness MemBind)
 -> TypeM MCMem (Param (MemInfo SubExp Uniqueness MemBind)))
-> Param (MemInfo SubExp Uniqueness MemBind)
-> TypeM MCMem (Param (MemInfo SubExp Uniqueness MemBind))
forall a b. (a -> b) -> a -> b
$ Attrs
-> VName
-> MemInfo SubExp Uniqueness MemBind
-> Param (MemInfo SubExp Uniqueness MemBind)
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
forall a. Monoid a => a
mempty VName
name (PrimType -> MemInfo SubExp Uniqueness MemBind
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
t)
  matchPat :: Pat (LetDec (Aliases MCMem))
-> Exp (Aliases MCMem) -> TypeM MCMem ()
matchPat = Pat (LetDec (Aliases MCMem))
-> Exp (Aliases MCMem) -> TypeM MCMem ()
forall rep inner.
(Mem rep inner, LetDec rep ~ MemBound NoUniqueness,
 Checkable rep) =>
Pat (LetDec (Aliases rep)) -> Exp (Aliases rep) -> TypeM rep ()
matchPatToExp
  matchReturnType :: [RetType MCMem] -> Result -> TypeM MCMem ()
matchReturnType = [RetType MCMem] -> Result -> TypeM MCMem ()
forall rep inner.
(Mem rep inner, Checkable rep) =>
[RetTypeMem] -> Result -> TypeM rep ()
matchFunctionReturnType
  matchBranchType :: [BranchType MCMem] -> Body (Aliases MCMem) -> TypeM MCMem ()
matchBranchType = [BranchType MCMem] -> Body (Aliases MCMem) -> TypeM MCMem ()
forall rep inner.
(Mem rep inner, Checkable rep) =>
[BodyReturns] -> Body (Aliases rep) -> TypeM rep ()
matchBranchReturnType
  matchLoopResult :: [FParam (Aliases MCMem)] -> Result -> TypeM MCMem ()
matchLoopResult = [FParam (Aliases MCMem)] -> Result -> TypeM MCMem ()
forall rep inner.
(Mem rep inner, Checkable rep) =>
[FParam (Aliases rep)] -> Result -> TypeM rep ()
matchLoopResultMem

instance BuilderOps MCMem where
  mkExpDecB :: Pat (LetDec MCMem) -> Exp MCMem -> m (ExpDec MCMem)
mkExpDecB Pat (LetDec MCMem)
_ Exp MCMem
_ = () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  mkBodyB :: Stms MCMem -> Result -> m (Body MCMem)
mkBodyB Stms MCMem
stms Result
res = Body MCMem -> m (Body MCMem)
forall (m :: * -> *) a. Monad m => a -> m a
return (Body MCMem -> m (Body MCMem)) -> Body MCMem -> m (Body MCMem)
forall a b. (a -> b) -> a -> b
$ BodyDec MCMem -> Stms MCMem -> Result -> Body MCMem
forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body () Stms MCMem
stms Result
res
  mkLetNamesB :: [VName] -> Exp MCMem -> m (Stm MCMem)
mkLetNamesB = ExpDec (Rep m) -> [VName] -> Exp (Rep m) -> m (Stm (Rep m))
forall (m :: * -> *) inner.
(LetDec (Rep m) ~ MemBound NoUniqueness, Mem (Rep m) inner,
 MonadBuilder m, ExpDec (Rep m) ~ ()) =>
ExpDec (Rep m) -> [VName] -> Exp (Rep m) -> m (Stm (Rep m))
mkLetNamesB' ()

instance BuilderOps (Engine.Wise MCMem) where
  mkExpDecB :: Pat (LetDec (Wise MCMem))
-> Exp (Wise MCMem) -> m (ExpDec (Wise MCMem))
mkExpDecB Pat (LetDec (Wise MCMem))
pat Exp (Wise MCMem)
e = (ExpWisdom, ()) -> m (ExpWisdom, ())
forall (m :: * -> *) a. Monad m => a -> m a
return ((ExpWisdom, ()) -> m (ExpWisdom, ()))
-> (ExpWisdom, ()) -> m (ExpWisdom, ())
forall a b. (a -> b) -> a -> b
$ Pat (LetDec (Wise MCMem))
-> ExpDec MCMem -> Exp (Wise MCMem) -> ExpDec (Wise MCMem)
forall rep.
(ASTRep rep, CanBeWise (Op rep)) =>
Pat (LetDec (Wise rep))
-> ExpDec rep -> Exp (Wise rep) -> ExpDec (Wise rep)
Engine.mkWiseExpDec Pat (LetDec (Wise MCMem))
pat () Exp (Wise MCMem)
e
  mkBodyB :: Stms (Wise MCMem) -> Result -> m (Body (Wise MCMem))
mkBodyB Stms (Wise MCMem)
stms Result
res = Body (Wise MCMem) -> m (Body (Wise MCMem))
forall (m :: * -> *) a. Monad m => a -> m a
return (Body (Wise MCMem) -> m (Body (Wise MCMem)))
-> Body (Wise MCMem) -> m (Body (Wise MCMem))
forall a b. (a -> b) -> a -> b
$ BodyDec MCMem -> Stms (Wise MCMem) -> Result -> Body (Wise MCMem)
forall rep.
(ASTRep rep, CanBeWise (Op rep)) =>
BodyDec rep -> Stms (Wise rep) -> Result -> Body (Wise rep)
Engine.mkWiseBody () Stms (Wise MCMem)
stms Result
res
  mkLetNamesB :: [VName] -> Exp (Wise MCMem) -> m (Stm (Wise MCMem))
mkLetNamesB = [VName] -> Exp (Wise MCMem) -> m (Stm (Wise MCMem))
forall rep inner (m :: * -> *).
(BuilderOps rep, Mem rep inner, LetDec rep ~ MemBound NoUniqueness,
 OpReturns (OpWithWisdom inner), ExpDec rep ~ (), Rep m ~ Wise rep,
 HasScope (Wise rep) m, MonadBuilder m, CanBeWise inner) =>
[VName] -> Exp (Wise rep) -> m (Stm (Wise rep))
mkLetNamesB''

instance TraverseOpStms (Engine.Wise MCMem) where
  traverseOpStms :: OpStmsTraverser m (Op (Wise MCMem)) (Wise MCMem)
traverseOpStms = OpStmsTraverser m (MCOp (Wise MCMem) ()) (Wise MCMem)
-> OpStmsTraverser m (MemOp (MCOp (Wise MCMem) ())) (Wise MCMem)
forall (m :: * -> *) inner rep.
Monad m =>
OpStmsTraverser m inner rep -> OpStmsTraverser m (MemOp inner) rep
traverseMemOpStms (OpStmsTraverser m () (Wise MCMem)
-> OpStmsTraverser m (MCOp (Wise MCMem) ()) (Wise MCMem)
forall (m :: * -> *) op rep.
Monad m =>
OpStmsTraverser m op rep -> OpStmsTraverser m (MCOp rep op) rep
traverseMCOpStms ((() -> m ()) -> OpStmsTraverser m () (Wise MCMem)
forall a b. a -> b -> a
const () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure))

simplifyProg :: Prog MCMem -> PassM (Prog MCMem)
simplifyProg :: Prog MCMem -> PassM (Prog MCMem)
simplifyProg = SimpleOps MCMem -> Prog MCMem -> PassM (Prog MCMem)
forall rep inner.
SimplifyMemory rep inner =>
SimpleOps rep -> Prog rep -> PassM (Prog rep)
simplifyProgGeneric SimpleOps MCMem
simpleMCMem

simpleMCMem :: Engine.SimpleOps MCMem
simpleMCMem :: SimpleOps MCMem
simpleMCMem =
  (OpWithWisdom (MCOp MCMem ()) -> UsageTable)
-> SimplifyOp MCMem (OpWithWisdom (MCOp MCMem ()))
-> SimpleOps MCMem
forall rep inner.
SimplifyMemory rep inner =>
(OpWithWisdom inner -> UsageTable)
-> SimplifyOp rep (OpWithWisdom inner) -> SimpleOps rep
simpleGeneric (UsageTable -> MCOp (Wise MCMem) () -> UsageTable
forall a b. a -> b -> a
const UsageTable
forall a. Monoid a => a
mempty) (SimplifyOp MCMem (OpWithWisdom (MCOp MCMem ()))
 -> SimpleOps MCMem)
-> SimplifyOp MCMem (OpWithWisdom (MCOp MCMem ()))
-> SimpleOps MCMem
forall a b. (a -> b) -> a -> b
$ SimplifyOp MCMem ()
-> MCOp (Wise MCMem) ()
-> SimpleM MCMem (MCOp (Wise MCMem) (), Stms (Wise MCMem))
forall rep op.
(SimplifiableRep rep, BodyDec rep ~ ()) =>
SimplifyOp rep op
-> MCOp (Wise rep) op
-> SimpleM rep (MCOp (Wise rep) op, Stms (Wise rep))
simplifyMCOp (SimplifyOp MCMem ()
 -> MCOp (Wise MCMem) ()
 -> SimpleM MCMem (MCOp (Wise MCMem) (), Stms (Wise MCMem)))
-> SimplifyOp MCMem ()
-> MCOp (Wise MCMem) ()
-> SimpleM MCMem (MCOp (Wise MCMem) (), Stms (Wise MCMem))
forall a b. (a -> b) -> a -> b
$ SimpleM MCMem ((), Stms (Wise MCMem)) -> SimplifyOp MCMem ()
forall a b. a -> b -> a
const (SimpleM MCMem ((), Stms (Wise MCMem)) -> SimplifyOp MCMem ())
-> SimpleM MCMem ((), Stms (Wise MCMem)) -> SimplifyOp MCMem ()
forall a b. (a -> b) -> a -> b
$ ((), Stms (Wise MCMem)) -> SimpleM MCMem ((), Stms (Wise MCMem))
forall (m :: * -> *) a. Monad m => a -> m a
return ((), Stms (Wise MCMem)
forall a. Monoid a => a
mempty)