{-# LANGUAGE OverloadedStrings #-}
module Futhark.CodeGen.ImpCode.Multicore
( Program,
Multicore (..),
MCCode,
Scheduling (..),
SchedulerInfo (..),
AtomicOp (..),
ParallelTask (..),
KernelHandling (..),
lexicalMemoryUsageMC,
module Futhark.CodeGen.ImpCode,
)
where
import qualified Data.Map as M
import Futhark.CodeGen.ImpCode
import Futhark.Util.Pretty
type Program = Functions Multicore
data Multicore
= SegOp String [Param] ParallelTask (Maybe ParallelTask) [Param] SchedulerInfo
| ParLoop String MCCode [Param]
|
ISPCKernel MCCode [Param]
|
ForEach VName Exp Exp MCCode
|
ForEachActive VName MCCode
|
VName Exp Exp
|
GetLoopBounds VName VName
|
GetTaskId VName
|
GetNumTasks VName
| Atomic AtomicOp
type MCCode = Code Multicore
data AtomicOp
= AtomicAdd IntType VName VName (Count Elements (TExp Int32)) Exp
| AtomicSub IntType VName VName (Count Elements (TExp Int32)) Exp
| AtomicAnd IntType VName VName (Count Elements (TExp Int32)) Exp
| AtomicOr IntType VName VName (Count Elements (TExp Int32)) Exp
| AtomicXor IntType VName VName (Count Elements (TExp Int32)) Exp
| AtomicXchg PrimType VName VName (Count Elements (TExp Int32)) Exp
| AtomicCmpXchg PrimType VName VName (Count Elements (TExp Int32)) VName Exp
deriving (Int -> AtomicOp -> ShowS
[AtomicOp] -> ShowS
AtomicOp -> String
(Int -> AtomicOp -> ShowS)
-> (AtomicOp -> String) -> ([AtomicOp] -> ShowS) -> Show AtomicOp
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [AtomicOp] -> ShowS
$cshowList :: [AtomicOp] -> ShowS
show :: AtomicOp -> String
$cshow :: AtomicOp -> String
showsPrec :: Int -> AtomicOp -> ShowS
$cshowsPrec :: Int -> AtomicOp -> ShowS
Show)
instance FreeIn AtomicOp where
freeIn' :: AtomicOp -> FV
freeIn' (AtomicAdd IntType
_ VName
_ VName
arr Count Elements (TExp Int32)
i Exp
x) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
arr FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Count Elements (TExp Int32) -> FV
forall a. FreeIn a => a -> FV
freeIn' Count Elements (TExp Int32)
i FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
x
freeIn' (AtomicSub IntType
_ VName
_ VName
arr Count Elements (TExp Int32)
i Exp
x) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
arr FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Count Elements (TExp Int32) -> FV
forall a. FreeIn a => a -> FV
freeIn' Count Elements (TExp Int32)
i FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
x
freeIn' (AtomicAnd IntType
_ VName
_ VName
arr Count Elements (TExp Int32)
i Exp
x) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
arr FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Count Elements (TExp Int32) -> FV
forall a. FreeIn a => a -> FV
freeIn' Count Elements (TExp Int32)
i FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
x
freeIn' (AtomicOr IntType
_ VName
_ VName
arr Count Elements (TExp Int32)
i Exp
x) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
arr FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Count Elements (TExp Int32) -> FV
forall a. FreeIn a => a -> FV
freeIn' Count Elements (TExp Int32)
i FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
x
freeIn' (AtomicXor IntType
_ VName
_ VName
arr Count Elements (TExp Int32)
i Exp
x) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
arr FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Count Elements (TExp Int32) -> FV
forall a. FreeIn a => a -> FV
freeIn' Count Elements (TExp Int32)
i FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
x
freeIn' (AtomicCmpXchg PrimType
_ VName
_ VName
arr Count Elements (TExp Int32)
i VName
retval Exp
x) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
arr FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Count Elements (TExp Int32) -> FV
forall a. FreeIn a => a -> FV
freeIn' Count Elements (TExp Int32)
i FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
x FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
retval
freeIn' (AtomicXchg PrimType
_ VName
_ VName
arr Count Elements (TExp Int32)
i Exp
x) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
arr FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Count Elements (TExp Int32) -> FV
forall a. FreeIn a => a -> FV
freeIn' Count Elements (TExp Int32)
i FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
x
data SchedulerInfo = SchedulerInfo
{
SchedulerInfo -> Exp
iterations :: Exp,
SchedulerInfo -> Scheduling
scheduling :: Scheduling
}
newtype ParallelTask = ParallelTask MCCode
data Scheduling
= Dynamic
| Static
instance Pretty Scheduling where
ppr :: Scheduling -> Doc
ppr Scheduling
Dynamic = Doc
"Dynamic"
ppr Scheduling
Static = Doc
"Static"
instance Pretty SchedulerInfo where
ppr :: SchedulerInfo -> Doc
ppr (SchedulerInfo Exp
i Scheduling
sched) =
[Doc] -> Doc
stack
[ String -> String -> Doc -> Doc
nestedBlock String
"scheduling {" String
"}" (Scheduling -> Doc
forall a. Pretty a => a -> Doc
ppr Scheduling
sched),
String -> String -> Doc -> Doc
nestedBlock String
"iter {" String
"}" (Exp -> Doc
forall a. Pretty a => a -> Doc
ppr Exp
i)
]
instance Pretty ParallelTask where
ppr :: ParallelTask -> Doc
ppr (ParallelTask MCCode
code) = MCCode -> Doc
forall a. Pretty a => a -> Doc
ppr MCCode
code
instance Pretty Multicore where
ppr :: Multicore -> Doc
ppr (GetLoopBounds VName
start VName
end) =
(VName, VName) -> Doc
forall a. Pretty a => a -> Doc
ppr (VName
start, VName
end) Doc -> Doc -> Doc
<+> Doc
"<-" Doc -> Doc -> Doc
<+> Doc
"get_loop_bounds()"
ppr (GetTaskId VName
v) =
VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
v Doc -> Doc -> Doc
<+> Doc
"<-" Doc -> Doc -> Doc
<+> Doc
"get_task_id()"
ppr (GetNumTasks VName
v) =
VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
v Doc -> Doc -> Doc
<+> Doc
"<-" Doc -> Doc -> Doc
<+> Doc
"get_num_tasks()"
ppr (SegOp String
s [Param]
free ParallelTask
seq_code Maybe ParallelTask
par_code [Param]
retval SchedulerInfo
scheduler) =
Doc
"SegOp" Doc -> Doc -> Doc
<+> String -> Doc
text String
s Doc -> Doc -> Doc
<+> String -> String -> Doc -> Doc
nestedBlock String
"{" String
"}" Doc
ppbody
where
ppbody :: Doc
ppbody =
[Doc] -> Doc
stack
[ SchedulerInfo -> Doc
forall a. Pretty a => a -> Doc
ppr SchedulerInfo
scheduler,
String -> String -> Doc -> Doc
nestedBlock String
"free {" String
"}" ([Param] -> Doc
forall a. Pretty a => a -> Doc
ppr [Param]
free),
String -> String -> Doc -> Doc
nestedBlock String
"seq {" String
"}" (ParallelTask -> Doc
forall a. Pretty a => a -> Doc
ppr ParallelTask
seq_code),
Doc -> (ParallelTask -> Doc) -> Maybe ParallelTask -> Doc
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Doc
forall a. Monoid a => a
mempty (String -> String -> Doc -> Doc
nestedBlock String
"par {" String
"}" (Doc -> Doc) -> (ParallelTask -> Doc) -> ParallelTask -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ParallelTask -> Doc
forall a. Pretty a => a -> Doc
ppr) Maybe ParallelTask
par_code,
String -> String -> Doc -> Doc
nestedBlock String
"retvals {" String
"}" ([Param] -> Doc
forall a. Pretty a => a -> Doc
ppr [Param]
retval)
]
ppr (ParLoop String
s MCCode
body [Param]
params) =
Doc
"parloop" Doc -> Doc -> Doc
<+> String -> Doc
forall a. Pretty a => a -> Doc
ppr String
s Doc -> Doc -> Doc
</> String -> String -> Doc -> Doc
nestedBlock String
"{" String
"}" Doc
ppbody
where
ppbody :: Doc
ppbody =
[Doc] -> Doc
stack
[ String -> String -> Doc -> Doc
nestedBlock String
"params {" String
"}" ([Param] -> Doc
forall a. Pretty a => a -> Doc
ppr [Param]
params),
String -> String -> Doc -> Doc
nestedBlock String
"body {" String
"}" (MCCode -> Doc
forall a. Pretty a => a -> Doc
ppr MCCode
body)
]
ppr (Atomic AtomicOp
_) =
Doc
"AtomicOp"
ppr (ISPCKernel MCCode
body [Param]
_) =
Doc
"ispc" Doc -> Doc -> Doc
<+> String -> String -> Doc -> Doc
nestedBlock String
"{" String
"}" (MCCode -> Doc
forall a. Pretty a => a -> Doc
ppr MCCode
body)
ppr (ForEach VName
i Exp
from Exp
to MCCode
body) =
Doc
"foreach"
Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
i
Doc -> Doc -> Doc
<+> Doc
"="
Doc -> Doc -> Doc
<+> Exp -> Doc
forall a. Pretty a => a -> Doc
ppr Exp
from
Doc -> Doc -> Doc
<+> Doc
"to"
Doc -> Doc -> Doc
<+> Exp -> Doc
forall a. Pretty a => a -> Doc
ppr Exp
to
Doc -> Doc -> Doc
<+> String -> String -> Doc -> Doc
nestedBlock String
"{" String
"}" (MCCode -> Doc
forall a. Pretty a => a -> Doc
ppr MCCode
body)
ppr (ForEachActive VName
i MCCode
body) =
Doc
"foreach_active"
Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
i
Doc -> Doc -> Doc
<+> String -> String -> Doc -> Doc
nestedBlock String
"{" String
"}" (MCCode -> Doc
forall a. Pretty a => a -> Doc
ppr MCCode
body)
ppr (ExtractLane VName
dest Exp
tar Exp
lane) =
VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
dest Doc -> Doc -> Doc
<+> Doc
"<-" Doc -> Doc -> Doc
<+> Doc
"extract" Doc -> Doc -> Doc
<+> Doc -> Doc
parens ([Doc] -> Doc
commasep ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (Exp -> Doc) -> [Exp] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Doc
forall a. Pretty a => a -> Doc
ppr [Exp
tar, Exp
lane])
instance FreeIn SchedulerInfo where
freeIn' :: SchedulerInfo -> FV
freeIn' (SchedulerInfo Exp
iter Scheduling
_) = Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
iter
instance FreeIn ParallelTask where
freeIn' :: ParallelTask -> FV
freeIn' (ParallelTask MCCode
code) = MCCode -> FV
forall a. FreeIn a => a -> FV
freeIn' MCCode
code
instance FreeIn Multicore where
freeIn' :: Multicore -> FV
freeIn' (GetLoopBounds VName
start VName
end) =
(VName, VName) -> FV
forall a. FreeIn a => a -> FV
freeIn' (VName
start, VName
end)
freeIn' (GetTaskId VName
v) =
VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
v
freeIn' (GetNumTasks VName
v) =
VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
v
freeIn' (SegOp String
_ [Param]
_ ParallelTask
par_code Maybe ParallelTask
seq_code [Param]
_ SchedulerInfo
info) =
ParallelTask -> FV
forall a. FreeIn a => a -> FV
freeIn' ParallelTask
par_code FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Maybe ParallelTask -> FV
forall a. FreeIn a => a -> FV
freeIn' Maybe ParallelTask
seq_code FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> SchedulerInfo -> FV
forall a. FreeIn a => a -> FV
freeIn' SchedulerInfo
info
freeIn' (ParLoop String
_ MCCode
body [Param]
_) =
MCCode -> FV
forall a. FreeIn a => a -> FV
freeIn' MCCode
body
freeIn' (Atomic AtomicOp
aop) =
AtomicOp -> FV
forall a. FreeIn a => a -> FV
freeIn' AtomicOp
aop
freeIn' (ISPCKernel MCCode
body [Param]
_) =
MCCode -> FV
forall a. FreeIn a => a -> FV
freeIn' MCCode
body
freeIn' (ForEach VName
i Exp
from Exp
to MCCode
body) =
Names -> FV -> FV
fvBind (VName -> Names
oneName VName
i) (MCCode -> FV
forall a. FreeIn a => a -> FV
freeIn' MCCode
body FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
from FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
to)
freeIn' (ForEachActive VName
i MCCode
body) =
Names -> FV -> FV
fvBind (VName -> Names
oneName VName
i) (MCCode -> FV
forall a. FreeIn a => a -> FV
freeIn' MCCode
body)
freeIn' (ExtractLane VName
dest Exp
tar Exp
lane) =
VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
dest FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
tar FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
lane
data KernelHandling = TraverseKernels | OpaqueKernels
lexicalMemoryUsageMC :: KernelHandling -> Function Multicore -> M.Map VName Space
lexicalMemoryUsageMC :: KernelHandling -> Function Multicore -> Map VName Space
lexicalMemoryUsageMC KernelHandling
gokernel Function Multicore
func =
(VName -> Space -> Bool) -> Map VName Space -> Map VName Space
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (Bool -> Space -> Bool
forall a b. a -> b -> a
const (Bool -> Space -> Bool)
-> (VName -> Bool) -> VName -> Space -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Names -> Bool
`notNameIn` Names
nonlexical)) (Map VName Space -> Map VName Space)
-> Map VName Space -> Map VName Space
forall a b. (a -> b) -> a -> b
$
MCCode -> Map VName Space
declared (MCCode -> Map VName Space) -> MCCode -> Map VName Space
forall a b. (a -> b) -> a -> b
$
Function Multicore -> MCCode
forall a. FunctionT a -> Code a
functionBody Function Multicore
func
where
nonlexical :: Names
nonlexical =
MCCode -> Names
set (Function Multicore -> MCCode
forall a. FunctionT a -> Code a
functionBody Function Multicore
func)
Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [VName] -> Names
namesFromList ((Param -> VName) -> [Param] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param -> VName
paramName (Function Multicore -> [Param]
forall a. FunctionT a -> [Param]
functionOutput Function Multicore
func))
go :: (MCCode -> a) -> MCCode -> a
go MCCode -> a
f (MCCode
x :>>: MCCode
y) = MCCode -> a
f MCCode
x a -> a -> a
forall a. Semigroup a => a -> a -> a
<> MCCode -> a
f MCCode
y
go MCCode -> a
f (If TExp Bool
_ MCCode
x MCCode
y) = MCCode -> a
f MCCode
x a -> a -> a
forall a. Semigroup a => a -> a -> a
<> MCCode -> a
f MCCode
y
go MCCode -> a
f (For VName
_ Exp
_ MCCode
x) = MCCode -> a
f MCCode
x
go MCCode -> a
f (While TExp Bool
_ MCCode
x) = MCCode -> a
f MCCode
x
go MCCode -> a
f (Comment String
_ MCCode
x) = MCCode -> a
f MCCode
x
go MCCode -> a
f (Op Multicore
op) = (MCCode -> a) -> Multicore -> a
goOp MCCode -> a
f Multicore
op
go MCCode -> a
_ MCCode
_ = a
forall a. Monoid a => a
mempty
goOp :: (MCCode -> a) -> Multicore -> a
goOp MCCode -> a
f (ForEach VName
_ Exp
_ Exp
_ MCCode
body) = (MCCode -> a) -> MCCode -> a
go MCCode -> a
f MCCode
body
goOp MCCode -> a
f (ForEachActive VName
_ MCCode
body) = (MCCode -> a) -> MCCode -> a
go MCCode -> a
f MCCode
body
goOp MCCode -> a
f (ISPCKernel MCCode
body [Param]
_) =
case KernelHandling
gokernel of
KernelHandling
TraverseKernels -> (MCCode -> a) -> MCCode -> a
go MCCode -> a
f MCCode
body
KernelHandling
OpaqueKernels -> a
forall a. Monoid a => a
mempty
goOp MCCode -> a
_ Multicore
_ = a
forall a. Monoid a => a
mempty
declared :: MCCode -> Map VName Space
declared (DeclareMem VName
mem Space
spc) =
VName -> Space -> Map VName Space
forall k a. k -> a -> Map k a
M.singleton VName
mem Space
spc
declared MCCode
x = (MCCode -> Map VName Space) -> MCCode -> Map VName Space
forall a. Monoid a => (MCCode -> a) -> MCCode -> a
go MCCode -> Map VName Space
declared MCCode
x
set :: MCCode -> Names
set (SetMem VName
x VName
y Space
_) = [VName] -> Names
namesFromList [VName
x, VName
y]
set (Call [VName]
_ Name
_ [Arg]
args) = (Arg -> Names) -> [Arg] -> Names
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Arg -> Names
onArg [Arg]
args
where
onArg :: Arg -> Names
onArg ExpArg {} = Names
forall a. Monoid a => a
mempty
onArg (MemArg VName
x) = VName -> Names
oneName VName
x
set (Op (SegOp String
_ [Param]
params ParallelTask
_ Maybe ParallelTask
_ [Param]
retvals SchedulerInfo
_)) =
case KernelHandling
gokernel of
KernelHandling
TraverseKernels -> Names
forall a. Monoid a => a
mempty
KernelHandling
OpaqueKernels -> [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ (Param -> VName) -> [Param] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param -> VName
paramName [Param]
params [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> (Param -> VName) -> [Param] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param -> VName
paramName [Param]
retvals
set MCCode
x = (MCCode -> Names) -> MCCode -> Names
forall a. Monoid a => (MCCode -> a) -> MCCode -> a
go MCCode -> Names
set MCCode
x