{-# LANGUAGE TypeFamilies #-}
module Futhark.IR.GPUMem
( GPUMem,
simplifyProg,
simplifyStms,
simpleGPUMem,
module Futhark.IR.Mem,
module Futhark.IR.GPU.Op,
)
where
import Futhark.Analysis.PrimExp.Convert
import Futhark.Analysis.UsageTable qualified as UT
import Futhark.IR.Aliases (Aliases)
import Futhark.IR.GPU.Op
import Futhark.IR.GPU.Simplify (simplifyKernelOp)
import Futhark.IR.Mem
import Futhark.IR.Mem.Simplify
import Futhark.IR.TypeCheck qualified as TC
import Futhark.MonadFreshNames
import Futhark.Optimise.Simplify.Engine qualified as Engine
import Futhark.Pass
import Futhark.Pass.ExplicitAllocations (BuilderOps (..), mkLetNamesB', mkLetNamesB'')
data GPUMem
instance RepTypes GPUMem where
type LetDec GPUMem = LetDecMem
type FParamInfo GPUMem = FParamMem
type LParamInfo GPUMem = LParamMem
type RetType GPUMem = RetTypeMem
type BranchType GPUMem = BranchTypeMem
type OpC GPUMem = MemOp (HostOp NoOp)
instance ASTRep GPUMem where
expTypesFromPat :: forall (m :: * -> *).
(HasScope GPUMem m, Monad m) =>
Pat (LetDec GPUMem) -> m [BranchType GPUMem]
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 GPUMem
instance TC.Checkable GPUMem where
checkOp :: Op (Aliases GPUMem) -> TypeM GPUMem ()
checkOp = Maybe SegLevel
-> MemOp (HostOp NoOp) (Aliases GPUMem) -> TypeM GPUMem ()
typeCheckMemoryOp Maybe SegLevel
forall a. Maybe a
Nothing
where
typeCheckMemoryOp ::
Maybe SegLevel ->
MemOp (HostOp NoOp) (Aliases GPUMem) ->
TC.TypeM GPUMem ()
typeCheckMemoryOp :: Maybe SegLevel
-> MemOp (HostOp NoOp) (Aliases GPUMem) -> TypeM GPUMem ()
typeCheckMemoryOp Maybe SegLevel
_ (Alloc SubExp
size Space
_) =
[TypeBase (ShapeBase SubExp) NoUniqueness]
-> SubExp -> TypeM GPUMem ()
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
typeCheckMemoryOp Maybe SegLevel
lvl (Inner HostOp NoOp (Aliases GPUMem)
op) =
(SegLevel -> Op (Aliases GPUMem) -> TypeM GPUMem ())
-> Maybe SegLevel
-> (NoOp (Aliases GPUMem) -> TypeM GPUMem ())
-> HostOp NoOp (Aliases GPUMem)
-> TypeM GPUMem ()
forall rep (op :: * -> *).
Checkable rep =>
(SegLevel -> Op (Aliases rep) -> TypeM rep ())
-> Maybe SegLevel
-> (op (Aliases rep) -> TypeM rep ())
-> HostOp op (Aliases rep)
-> TypeM rep ()
typeCheckHostOp (Maybe SegLevel
-> MemOp (HostOp NoOp) (Aliases GPUMem) -> TypeM GPUMem ()
typeCheckMemoryOp (Maybe SegLevel
-> MemOp (HostOp NoOp) (Aliases GPUMem) -> TypeM GPUMem ())
-> (SegLevel -> Maybe SegLevel)
-> SegLevel
-> MemOp (HostOp NoOp) (Aliases GPUMem)
-> TypeM GPUMem ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SegLevel -> Maybe SegLevel
forall a. a -> Maybe a
Just) Maybe SegLevel
lvl (TypeM GPUMem () -> NoOp (Aliases GPUMem) -> TypeM GPUMem ()
forall a b. a -> b -> a
const (TypeM GPUMem () -> NoOp (Aliases GPUMem) -> TypeM GPUMem ())
-> TypeM GPUMem () -> NoOp (Aliases GPUMem) -> TypeM GPUMem ()
forall a b. (a -> b) -> a -> b
$ () -> TypeM GPUMem ()
forall a. a -> TypeM GPUMem a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) HostOp NoOp (Aliases GPUMem)
op
checkFParamDec :: VName -> FParamInfo GPUMem -> TypeM GPUMem ()
checkFParamDec = VName -> FParamInfo GPUMem -> TypeM GPUMem ()
VName -> FParamMem -> TypeM GPUMem ()
forall rep u.
Checkable rep =>
VName -> MemInfo SubExp u MemBind -> TypeM rep ()
checkMemInfo
checkLParamDec :: VName -> LParamInfo GPUMem -> TypeM GPUMem ()
checkLParamDec = VName -> LParamInfo GPUMem -> TypeM GPUMem ()
VName -> LetDecMem -> TypeM GPUMem ()
forall rep u.
Checkable rep =>
VName -> MemInfo SubExp u MemBind -> TypeM rep ()
checkMemInfo
checkLetBoundDec :: VName -> LetDec GPUMem -> TypeM GPUMem ()
checkLetBoundDec = VName -> LetDec GPUMem -> TypeM GPUMem ()
VName -> LetDecMem -> TypeM GPUMem ()
forall rep u.
Checkable rep =>
VName -> MemInfo SubExp u MemBind -> TypeM rep ()
checkMemInfo
checkRetType :: [RetType GPUMem] -> TypeM GPUMem ()
checkRetType = (RetType GPUMem -> TypeM GPUMem ())
-> [RetType GPUMem] -> TypeM GPUMem ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((RetType GPUMem -> TypeM GPUMem ())
-> [RetType GPUMem] -> TypeM GPUMem ())
-> (RetType GPUMem -> TypeM GPUMem ())
-> [RetType GPUMem]
-> TypeM GPUMem ()
forall a b. (a -> b) -> a -> b
$ TypeBase (ShapeBase (Ext SubExp)) Uniqueness -> TypeM GPUMem ()
forall rep u.
Checkable rep =>
TypeBase (ShapeBase (Ext SubExp)) u -> TypeM rep ()
TC.checkExtType (TypeBase (ShapeBase (Ext SubExp)) Uniqueness -> TypeM GPUMem ())
-> (RetTypeMem -> TypeBase (ShapeBase (Ext SubExp)) Uniqueness)
-> RetTypeMem
-> TypeM GPUMem ()
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 GPUMem (FParam (Aliases GPUMem))
primFParam VName
name PrimType
t = FParam (Aliases GPUMem) -> TypeM GPUMem (FParam (Aliases GPUMem))
forall a. a -> TypeM GPUMem a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FParam (Aliases GPUMem) -> TypeM GPUMem (FParam (Aliases GPUMem)))
-> FParam (Aliases GPUMem)
-> TypeM GPUMem (FParam (Aliases GPUMem))
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 GPUMem))
-> Exp (Aliases GPUMem) -> TypeM GPUMem ()
matchPat = Pat (LetDec (Aliases GPUMem))
-> Exp (Aliases GPUMem) -> TypeM GPUMem ()
forall rep (inner :: * -> *).
(Mem rep inner, LetDec rep ~ LetDecMem, Checkable rep) =>
Pat (LetDec (Aliases rep)) -> Exp (Aliases rep) -> TypeM rep ()
matchPatToExp
matchReturnType :: [RetType GPUMem] -> Result -> TypeM GPUMem ()
matchReturnType = [RetType GPUMem] -> Result -> TypeM GPUMem ()
[RetTypeMem] -> Result -> TypeM GPUMem ()
forall rep (inner :: * -> *).
(Mem rep inner, Checkable rep) =>
[RetTypeMem] -> Result -> TypeM rep ()
matchFunctionReturnType
matchBranchType :: [BranchType GPUMem] -> Body (Aliases GPUMem) -> TypeM GPUMem ()
matchBranchType = [BranchType GPUMem] -> Body (Aliases GPUMem) -> TypeM GPUMem ()
[BranchTypeMem] -> Body (Aliases GPUMem) -> TypeM GPUMem ()
forall rep (inner :: * -> *).
(Mem rep inner, Checkable rep) =>
[BranchTypeMem] -> Body (Aliases rep) -> TypeM rep ()
matchBranchReturnType
matchLoopResult :: [FParam (Aliases GPUMem)] -> Result -> TypeM GPUMem ()
matchLoopResult = [FParam (Aliases GPUMem)] -> Result -> TypeM GPUMem ()
forall rep (inner :: * -> *).
(Mem rep inner, Checkable rep) =>
[FParam (Aliases rep)] -> Result -> TypeM rep ()
matchLoopResultMem
instance BuilderOps GPUMem where
mkExpDecB :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPUMem) =>
Pat (LetDec GPUMem) -> Exp GPUMem -> m (ExpDec GPUMem)
mkExpDecB Pat (LetDec GPUMem)
_ Exp GPUMem
_ = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
mkBodyB :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPUMem) =>
Stms GPUMem -> Result -> m (Body GPUMem)
mkBodyB Stms GPUMem
stms Result
res = Body GPUMem -> m (Body GPUMem)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Body GPUMem -> m (Body GPUMem)) -> Body GPUMem -> m (Body GPUMem)
forall a b. (a -> b) -> a -> b
$ BodyDec GPUMem -> Stms GPUMem -> Result -> Body GPUMem
forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body () Stms GPUMem
stms Result
res
mkLetNamesB :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPUMem) =>
[VName] -> Exp GPUMem -> m (Stm GPUMem)
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' (SpaceId -> Space
Space SpaceId
"device") ()
instance BuilderOps (Engine.Wise GPUMem) where
mkExpDecB :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ Wise GPUMem) =>
Pat (LetDec (Wise GPUMem))
-> Exp (Wise GPUMem) -> m (ExpDec (Wise GPUMem))
mkExpDecB Pat (LetDec (Wise GPUMem))
pat Exp (Wise GPUMem)
e = ExpDec (Wise GPUMem) -> m (ExpDec (Wise GPUMem))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpDec (Wise GPUMem) -> m (ExpDec (Wise GPUMem)))
-> ExpDec (Wise GPUMem) -> m (ExpDec (Wise GPUMem))
forall a b. (a -> b) -> a -> b
$ Pat (LetDec (Wise GPUMem))
-> ExpDec GPUMem -> Exp (Wise GPUMem) -> ExpDec (Wise GPUMem)
forall rep.
Informing rep =>
Pat (LetDec (Wise rep))
-> ExpDec rep -> Exp (Wise rep) -> ExpDec (Wise rep)
Engine.mkWiseExpDec Pat (LetDec (Wise GPUMem))
pat () Exp (Wise GPUMem)
e
mkBodyB :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ Wise GPUMem) =>
Stms (Wise GPUMem) -> Result -> m (Body (Wise GPUMem))
mkBodyB Stms (Wise GPUMem)
stms Result
res = Body (Wise GPUMem) -> m (Body (Wise GPUMem))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Body (Wise GPUMem) -> m (Body (Wise GPUMem)))
-> Body (Wise GPUMem) -> m (Body (Wise GPUMem))
forall a b. (a -> b) -> a -> b
$ BodyDec GPUMem
-> Stms (Wise GPUMem) -> Result -> Body (Wise GPUMem)
forall rep.
Informing rep =>
BodyDec rep -> Stms (Wise rep) -> Result -> Body (Wise rep)
Engine.mkWiseBody () Stms (Wise GPUMem)
stms Result
res
mkLetNamesB :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ Wise GPUMem) =>
[VName] -> Exp (Wise GPUMem) -> m (Stm (Wise GPUMem))
mkLetNamesB = Space -> [VName] -> Exp (Wise GPUMem) -> m (Stm (Wise GPUMem))
forall rep (inner :: * -> *) (m :: * -> *).
(Mem rep inner, LetDec rep ~ LetDecMem, OpReturns inner,
ExpDec rep ~ (), Rep m ~ Wise rep, HasScope (Wise rep) m,
MonadBuilder m, AliasedOp inner, RephraseOp (MemOp inner),
CanBeWise inner, ASTConstraints (inner (Wise rep))) =>
Space -> [VName] -> Exp (Wise rep) -> m (Stm (Wise rep))
mkLetNamesB'' (SpaceId -> Space
Space SpaceId
"device")
instance TraverseOpStms (Engine.Wise GPUMem) where
traverseOpStms :: forall (m :: * -> *).
Monad m =>
OpStmsTraverser m (Op (Wise GPUMem)) (Wise GPUMem)
traverseOpStms = OpStmsTraverser m (HostOp NoOp (Wise GPUMem)) (Wise GPUMem)
-> OpStmsTraverser
m (MemOp (HostOp NoOp) (Wise GPUMem)) (Wise GPUMem)
forall (m :: * -> *) (inner :: * -> *) rep.
Monad m =>
OpStmsTraverser m (inner rep) rep
-> OpStmsTraverser m (MemOp inner rep) rep
traverseMemOpStms (OpStmsTraverser m (NoOp (Wise GPUMem)) (Wise GPUMem)
-> OpStmsTraverser m (HostOp NoOp (Wise GPUMem)) (Wise GPUMem)
forall (m :: * -> *) (op :: * -> *) rep.
Monad m =>
OpStmsTraverser m (op rep) rep
-> OpStmsTraverser m (HostOp op rep) rep
traverseHostOpStms ((NoOp (Wise GPUMem) -> m (NoOp (Wise GPUMem)))
-> OpStmsTraverser m (NoOp (Wise GPUMem)) (Wise GPUMem)
forall a b. a -> b -> a
const NoOp (Wise GPUMem) -> m (NoOp (Wise GPUMem))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure))
simplifyProg :: Prog GPUMem -> PassM (Prog GPUMem)
simplifyProg :: Prog GPUMem -> PassM (Prog GPUMem)
simplifyProg = RuleBook (Wise GPUMem)
-> SimpleOps GPUMem -> Prog GPUMem -> PassM (Prog GPUMem)
forall rep (inner :: * -> *).
SimplifyMemory rep inner =>
RuleBook (Wise rep)
-> SimpleOps rep -> Prog rep -> PassM (Prog rep)
simplifyProgGeneric RuleBook (Wise GPUMem)
forall rep (inner :: * -> *).
SimplifyMemory rep inner =>
RuleBook (Wise rep)
memRuleBook SimpleOps GPUMem
simpleGPUMem
simplifyStms ::
(HasScope GPUMem m, MonadFreshNames m) => Stms GPUMem -> m (Stms GPUMem)
simplifyStms :: forall (m :: * -> *).
(HasScope GPUMem m, MonadFreshNames m) =>
Stms GPUMem -> m (Stms GPUMem)
simplifyStms = RuleBook (Wise GPUMem)
-> SimpleOps GPUMem -> Stms GPUMem -> m (Stms GPUMem)
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, MonadFreshNames m, SimplifyMemory rep inner) =>
RuleBook (Wise rep) -> SimpleOps rep -> Stms rep -> m (Stms rep)
simplifyStmsGeneric RuleBook (Wise GPUMem)
forall rep (inner :: * -> *).
SimplifyMemory rep inner =>
RuleBook (Wise rep)
memRuleBook SimpleOps GPUMem
simpleGPUMem
simpleGPUMem :: Engine.SimpleOps GPUMem
simpleGPUMem :: SimpleOps GPUMem
simpleGPUMem =
(HostOp NoOp (Wise GPUMem) -> UsageTable)
-> SimplifyOp GPUMem (HostOp NoOp (Wise GPUMem))
-> SimpleOps GPUMem
forall rep (inner :: * -> *).
SimplifyMemory rep inner =>
(inner (Wise rep) -> UsageTable)
-> SimplifyOp rep (inner (Wise rep)) -> SimpleOps rep
simpleGeneric HostOp NoOp (Wise GPUMem) -> UsageTable
usage (SimplifyOp GPUMem (HostOp NoOp (Wise GPUMem)) -> SimpleOps GPUMem)
-> SimplifyOp GPUMem (HostOp NoOp (Wise GPUMem))
-> SimpleOps GPUMem
forall a b. (a -> b) -> a -> b
$ SimplifyOp GPUMem (NoOp (Wise GPUMem))
-> SimplifyOp GPUMem (HostOp NoOp (Wise GPUMem))
forall rep (op :: * -> *).
(SimplifiableRep rep, BodyDec rep ~ ()) =>
SimplifyOp rep (op (Wise rep))
-> HostOp op (Wise rep)
-> SimpleM rep (HostOp op (Wise rep), Stms (Wise rep))
simplifyKernelOp (SimplifyOp GPUMem (NoOp (Wise GPUMem))
-> SimplifyOp GPUMem (HostOp NoOp (Wise GPUMem)))
-> SimplifyOp GPUMem (NoOp (Wise GPUMem))
-> SimplifyOp GPUMem (HostOp NoOp (Wise GPUMem))
forall a b. (a -> b) -> a -> b
$ SimpleM GPUMem (NoOp (Wise GPUMem), Stms (Wise GPUMem))
-> SimplifyOp GPUMem (NoOp (Wise GPUMem))
forall a b. a -> b -> a
const (SimpleM GPUMem (NoOp (Wise GPUMem), Stms (Wise GPUMem))
-> SimplifyOp GPUMem (NoOp (Wise GPUMem)))
-> SimpleM GPUMem (NoOp (Wise GPUMem), Stms (Wise GPUMem))
-> SimplifyOp GPUMem (NoOp (Wise GPUMem))
forall a b. (a -> b) -> a -> b
$ (NoOp (Wise GPUMem), Stms (Wise GPUMem))
-> SimpleM GPUMem (NoOp (Wise GPUMem), Stms (Wise GPUMem))
forall a. a -> SimpleM GPUMem a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NoOp (Wise GPUMem)
forall {k} (rep :: k). NoOp rep
NoOp, Stms (Wise GPUMem)
forall a. Monoid a => a
mempty)
where
usage :: HostOp NoOp (Wise GPUMem) -> UsageTable
usage (SegOp (SegMap SegLevel
_ SegSpace
_ [TypeBase (ShapeBase SubExp) NoUniqueness]
_ KernelBody (Wise GPUMem)
kbody)) = KernelBody (Wise GPUMem) -> UsageTable
localAllocs KernelBody (Wise GPUMem)
kbody
usage HostOp NoOp (Wise GPUMem)
_ = UsageTable
forall a. Monoid a => a
mempty
localAllocs :: KernelBody (Wise GPUMem) -> UsageTable
localAllocs = (Stm (Wise GPUMem) -> UsageTable)
-> Stms (Wise GPUMem) -> UsageTable
forall m a. Monoid m => (a -> m) -> Seq a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm (Wise GPUMem) -> UsageTable
stmSharedAlloc (Stms (Wise GPUMem) -> UsageTable)
-> (KernelBody (Wise GPUMem) -> Stms (Wise GPUMem))
-> KernelBody (Wise GPUMem)
-> UsageTable
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KernelBody (Wise GPUMem) -> Stms (Wise GPUMem)
forall rep. KernelBody rep -> Stms rep
kernelBodyStms
stmSharedAlloc :: Stm (Wise GPUMem) -> UsageTable
stmSharedAlloc = Exp (Wise GPUMem) -> UsageTable
expSharedAlloc (Exp (Wise GPUMem) -> UsageTable)
-> (Stm (Wise GPUMem) -> Exp (Wise GPUMem))
-> Stm (Wise GPUMem)
-> UsageTable
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm (Wise GPUMem) -> Exp (Wise GPUMem)
forall rep. Stm rep -> Exp rep
stmExp
expSharedAlloc :: Exp (Wise GPUMem) -> UsageTable
expSharedAlloc (Op (Alloc (Var VName
v) Space
_)) =
VName -> UsageTable
UT.sizeUsage VName
v
expSharedAlloc (Op (Inner (SegOp (SegMap SegLevel
_ SegSpace
_ [TypeBase (ShapeBase SubExp) NoUniqueness]
_ KernelBody (Wise GPUMem)
kbody)))) =
KernelBody (Wise GPUMem) -> UsageTable
localAllocs KernelBody (Wise GPUMem)
kbody
expSharedAlloc Exp (Wise GPUMem)
_ =
UsageTable
forall a. Monoid a => a
mempty