module Futhark.CodeGen.ImpGen.Multicore.Base
( extractAllocations,
compileThreadResult,
Locks (..),
HostEnv (..),
AtomicBinOp,
MulticoreGen,
decideScheduling,
decideScheduling',
groupResultArrays,
renameSegBinOp,
freeParams,
renameHistOpLambda,
atomicUpdateLocking,
AtomicUpdate (..),
DoAtomicUpdate,
Locking (..),
getSpace,
getLoopBounds,
getIterationDomain,
getReturnParams,
segOpString,
ChunkLoopVectorization (..),
generateChunkLoop,
generateUniformizeLoop,
extractVectorLane,
inISPC,
toParam,
sLoopNestVectorized,
)
where
import Control.Monad
import Data.Bifunctor
import qualified Data.Map as M
import Data.Maybe
import qualified Futhark.CodeGen.ImpCode.Multicore as Imp
import Futhark.CodeGen.ImpGen
import Futhark.Error
import Futhark.IR.MCMem
import Futhark.MonadFreshNames
import Futhark.Transform.Rename
import Prelude hiding (quot, rem)
type AtomicBinOp =
BinOp ->
Maybe (VName -> VName -> Imp.Count Imp.Elements (Imp.TExp Int32) -> Imp.Exp -> Imp.AtomicOp)
data Locks = Locks
{ Locks -> VName
locksArray :: VName,
Locks -> Int
locksCount :: Int
}
data HostEnv = HostEnv
{ HostEnv -> AtomicBinOp
hostAtomics :: AtomicBinOp,
HostEnv -> Map VName Locks
hostLocks :: M.Map VName Locks
}
type MulticoreGen = ImpM MCMem HostEnv Imp.Multicore
segOpString :: SegOp () MCMem -> MulticoreGen String
segOpString :: SegOp () MCMem -> MulticoreGen String
segOpString SegMap {} = String -> MulticoreGen String
forall (f :: * -> *) a. Applicative f => a -> f a
pure String
"segmap"
segOpString SegRed {} = String -> MulticoreGen String
forall (f :: * -> *) a. Applicative f => a -> f a
pure String
"segred"
segOpString SegScan {} = String -> MulticoreGen String
forall (f :: * -> *) a. Applicative f => a -> f a
pure String
"segscan"
segOpString SegHist {} = String -> MulticoreGen String
forall (f :: * -> *) a. Applicative f => a -> f a
pure String
"seghist"
arrParam :: VName -> MulticoreGen Imp.Param
arrParam :: VName -> MulticoreGen Param
arrParam VName
arr = do
VarEntry MCMem
name_entry <- VName -> ImpM MCMem HostEnv Multicore (VarEntry MCMem)
forall rep r op. VName -> ImpM rep r op (VarEntry rep)
lookupVar VName
arr
case VarEntry MCMem
name_entry of
ArrayVar Maybe (Exp MCMem)
_ (ArrayEntry (MemLoc VName
mem [SubExp]
_ IxFun (TPrimExp Int64 VName)
_) PrimType
_) ->
Param -> MulticoreGen Param
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Param -> MulticoreGen Param) -> Param -> MulticoreGen Param
forall a b. (a -> b) -> a -> b
$ VName -> Space -> Param
Imp.MemParam VName
mem Space
DefaultSpace
VarEntry MCMem
_ -> String -> MulticoreGen Param
forall a. HasCallStack => String -> a
error (String -> MulticoreGen Param) -> String -> MulticoreGen Param
forall a b. (a -> b) -> a -> b
$ String
"arrParam: could not handle array " String -> String -> String
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Show a => a -> String
show VName
arr
toParam :: VName -> TypeBase shape u -> MulticoreGen [Imp.Param]
toParam :: forall shape u. VName -> TypeBase shape u -> MulticoreGen [Param]
toParam VName
name (Prim PrimType
pt) = [Param] -> MulticoreGen [Param]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName -> PrimType -> Param
Imp.ScalarParam VName
name PrimType
pt]
toParam VName
name (Mem Space
space) = [Param] -> MulticoreGen [Param]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName -> Space -> Param
Imp.MemParam VName
name Space
space]
toParam VName
name Array {} = Param -> [Param]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Param -> [Param]) -> MulticoreGen Param -> MulticoreGen [Param]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> MulticoreGen Param
arrParam VName
name
toParam VName
_name Acc {} = [Param] -> MulticoreGen [Param]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
getSpace :: SegOp () MCMem -> SegSpace
getSpace :: SegOp () MCMem -> SegSpace
getSpace (SegHist ()
_ SegSpace
space [HistOp MCMem]
_ [Type]
_ KernelBody MCMem
_) = SegSpace
space
getSpace (SegRed ()
_ SegSpace
space [SegBinOp MCMem]
_ [Type]
_ KernelBody MCMem
_) = SegSpace
space
getSpace (SegScan ()
_ SegSpace
space [SegBinOp MCMem]
_ [Type]
_ KernelBody MCMem
_) = SegSpace
space
getSpace (SegMap ()
_ SegSpace
space [Type]
_ KernelBody MCMem
_) = SegSpace
space
getLoopBounds :: MulticoreGen (Imp.TExp Int64, Imp.TExp Int64)
getLoopBounds :: MulticoreGen (TPrimExp Int64 VName, TPrimExp Int64 VName)
getLoopBounds = do
TV Int64
start <- String -> PrimType -> ImpM MCMem HostEnv Multicore (TV Int64)
forall rep r op t. String -> PrimType -> ImpM rep r op (TV t)
dPrim String
"start" PrimType
int64
TV Int64
end <- String -> PrimType -> ImpM MCMem HostEnv Multicore (TV Int64)
forall rep r op t. String -> PrimType -> ImpM rep r op (TV t)
dPrim String
"end" PrimType
int64
Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall op rep r. Code op -> ImpM rep r op ()
emit (Code Multicore -> ImpM MCMem HostEnv Multicore ())
-> Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ Multicore -> Code Multicore
forall a. a -> Code a
Imp.Op (Multicore -> Code Multicore) -> Multicore -> Code Multicore
forall a b. (a -> b) -> a -> b
$ VName -> VName -> Multicore
Imp.GetLoopBounds (TV Int64 -> VName
forall t. TV t -> VName
tvVar TV Int64
start) (TV Int64 -> VName
forall t. TV t -> VName
tvVar TV Int64
end)
(TPrimExp Int64 VName, TPrimExp Int64 VName)
-> MulticoreGen (TPrimExp Int64 VName, TPrimExp Int64 VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TV Int64 -> TPrimExp Int64 VName
forall t. TV t -> TExp t
tvExp TV Int64
start, TV Int64 -> TPrimExp Int64 VName
forall t. TV t -> TExp t
tvExp TV Int64
end)
getIterationDomain :: SegOp () MCMem -> SegSpace -> MulticoreGen (Imp.TExp Int64)
getIterationDomain :: SegOp () MCMem -> SegSpace -> MulticoreGen (TPrimExp Int64 VName)
getIterationDomain SegMap {} SegSpace
space = do
let ns :: [SubExp]
ns = ((VName, SubExp) -> SubExp) -> [(VName, SubExp)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> SubExp
forall a b. (a, b) -> b
snd ([(VName, SubExp)] -> [SubExp]) -> [(VName, SubExp)] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space
ns_64 :: [TPrimExp Int64 VName]
ns_64 = (SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
ns
TPrimExp Int64 VName -> MulticoreGen (TPrimExp Int64 VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TPrimExp Int64 VName -> MulticoreGen (TPrimExp Int64 VName))
-> TPrimExp Int64 VName -> MulticoreGen (TPrimExp Int64 VName)
forall a b. (a -> b) -> a -> b
$ [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [TPrimExp Int64 VName]
ns_64
getIterationDomain SegOp () MCMem
_ SegSpace
space = do
let ns :: [SubExp]
ns = ((VName, SubExp) -> SubExp) -> [(VName, SubExp)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> SubExp
forall a b. (a, b) -> b
snd ([(VName, SubExp)] -> [SubExp]) -> [(VName, SubExp)] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space
ns_64 :: [TPrimExp Int64 VName]
ns_64 = (SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
ns
case SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space of
[(VName, SubExp)
_] -> TPrimExp Int64 VName -> MulticoreGen (TPrimExp Int64 VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TPrimExp Int64 VName -> MulticoreGen (TPrimExp Int64 VName))
-> TPrimExp Int64 VName -> MulticoreGen (TPrimExp Int64 VName)
forall a b. (a -> b) -> a -> b
$ [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [TPrimExp Int64 VName]
ns_64
[(VName, SubExp)]
_ -> TPrimExp Int64 VName -> MulticoreGen (TPrimExp Int64 VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TPrimExp Int64 VName -> MulticoreGen (TPrimExp Int64 VName))
-> TPrimExp Int64 VName -> MulticoreGen (TPrimExp Int64 VName)
forall a b. (a -> b) -> a -> b
$ [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([TPrimExp Int64 VName] -> TPrimExp Int64 VName)
-> [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a b. (a -> b) -> a -> b
$ [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. [a] -> [a]
init [TPrimExp Int64 VName]
ns_64
getReturnParams :: Pat LetDecMem -> SegOp () MCMem -> MulticoreGen [Imp.Param]
getReturnParams :: Pat LetDecMem -> SegOp () MCMem -> MulticoreGen [Param]
getReturnParams Pat LetDecMem
pat SegRed {} =
([[Param]] -> [Param])
-> ImpM MCMem HostEnv Multicore [[Param]] -> MulticoreGen [Param]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[Param]] -> [Param]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (ImpM MCMem HostEnv Multicore [[Param]] -> MulticoreGen [Param])
-> ((PatElem LetDecMem -> MulticoreGen [Param])
-> ImpM MCMem HostEnv Multicore [[Param]])
-> (PatElem LetDecMem -> MulticoreGen [Param])
-> MulticoreGen [Param]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [PatElem LetDecMem]
-> (PatElem LetDecMem -> MulticoreGen [Param])
-> ImpM MCMem HostEnv Multicore [[Param]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (Pat LetDecMem -> [PatElem LetDecMem]
forall dec. Pat dec -> [PatElem dec]
patElems Pat LetDecMem
pat) ((PatElem LetDecMem -> MulticoreGen [Param])
-> MulticoreGen [Param])
-> (PatElem LetDecMem -> MulticoreGen [Param])
-> MulticoreGen [Param]
forall a b. (a -> b) -> a -> b
$ \PatElem LetDecMem
pe -> do
case PatElem LetDecMem -> Type
forall dec. Typed dec => PatElem dec -> Type
patElemType PatElem LetDecMem
pe of
Prim PrimType
pt -> PatElem LetDecMem -> VName
forall dec. PatElem dec -> VName
patElemName PatElem LetDecMem
pe VName -> Exp -> ImpM MCMem HostEnv Multicore ()
forall rep r op. VName -> Exp -> ImpM rep r op ()
<~~ PrimValue -> Exp
forall v. PrimValue -> PrimExp v
ValueExp (PrimType -> PrimValue
blankPrimValue PrimType
pt)
Type
_ -> () -> ImpM MCMem HostEnv Multicore ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
VName -> Type -> MulticoreGen [Param]
forall shape u. VName -> TypeBase shape u -> MulticoreGen [Param]
toParam (PatElem LetDecMem -> VName
forall dec. PatElem dec -> VName
patElemName PatElem LetDecMem
pe) (PatElem LetDecMem -> Type
forall dec. Typed dec => PatElem dec -> Type
patElemType PatElem LetDecMem
pe)
getReturnParams Pat LetDecMem
_ SegOp () MCMem
_ = [Param] -> MulticoreGen [Param]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Param]
forall a. Monoid a => a
mempty
renameSegBinOp :: [SegBinOp MCMem] -> MulticoreGen [SegBinOp MCMem]
renameSegBinOp :: [SegBinOp MCMem] -> MulticoreGen [SegBinOp MCMem]
renameSegBinOp [SegBinOp MCMem]
segbinops =
[SegBinOp MCMem]
-> (SegBinOp MCMem
-> ImpM MCMem HostEnv Multicore (SegBinOp MCMem))
-> MulticoreGen [SegBinOp MCMem]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [SegBinOp MCMem]
segbinops ((SegBinOp MCMem -> ImpM MCMem HostEnv Multicore (SegBinOp MCMem))
-> MulticoreGen [SegBinOp MCMem])
-> (SegBinOp MCMem
-> ImpM MCMem HostEnv Multicore (SegBinOp MCMem))
-> MulticoreGen [SegBinOp MCMem]
forall a b. (a -> b) -> a -> b
$ \(SegBinOp Commutativity
comm Lambda MCMem
lam [SubExp]
ne Shape
shape) -> do
Lambda MCMem
lam' <- Lambda MCMem -> ImpM MCMem HostEnv Multicore (Lambda MCMem)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Lambda rep -> m (Lambda rep)
renameLambda Lambda MCMem
lam
SegBinOp MCMem -> ImpM MCMem HostEnv Multicore (SegBinOp MCMem)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SegBinOp MCMem -> ImpM MCMem HostEnv Multicore (SegBinOp MCMem))
-> SegBinOp MCMem -> ImpM MCMem HostEnv Multicore (SegBinOp MCMem)
forall a b. (a -> b) -> a -> b
$ Commutativity
-> Lambda MCMem -> [SubExp] -> Shape -> SegBinOp MCMem
forall rep.
Commutativity -> Lambda rep -> [SubExp] -> Shape -> SegBinOp rep
SegBinOp Commutativity
comm Lambda MCMem
lam' [SubExp]
ne Shape
shape
compileThreadResult ::
SegSpace ->
PatElem LetDecMem ->
KernelResult ->
MulticoreGen ()
compileThreadResult :: SegSpace
-> PatElem LetDecMem
-> KernelResult
-> ImpM MCMem HostEnv Multicore ()
compileThreadResult SegSpace
space PatElem LetDecMem
pe (Returns ResultManifest
_ Certs
_ SubExp
what) = do
let is :: [TPrimExp Int64 VName]
is = ((VName, SubExp) -> TPrimExp Int64 VName)
-> [(VName, SubExp)] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 (VName -> TPrimExp Int64 VName)
-> ((VName, SubExp) -> VName)
-> (VName, SubExp)
-> TPrimExp Int64 VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, SubExp) -> VName
forall a b. (a, b) -> a
fst) ([(VName, SubExp)] -> [TPrimExp Int64 VName])
-> [(VName, SubExp)] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM MCMem HostEnv Multicore ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (PatElem LetDecMem -> VName
forall dec. PatElem dec -> VName
patElemName PatElem LetDecMem
pe) [TPrimExp Int64 VName]
is SubExp
what []
compileThreadResult SegSpace
_ PatElem LetDecMem
_ ConcatReturns {} =
String -> ImpM MCMem HostEnv Multicore ()
forall a. String -> a
compilerBugS String
"compileThreadResult: ConcatReturn unhandled."
compileThreadResult SegSpace
_ PatElem LetDecMem
_ WriteReturns {} =
String -> ImpM MCMem HostEnv Multicore ()
forall a. String -> a
compilerBugS String
"compileThreadResult: WriteReturns unhandled."
compileThreadResult SegSpace
_ PatElem LetDecMem
_ TileReturns {} =
String -> ImpM MCMem HostEnv Multicore ()
forall a. String -> a
compilerBugS String
"compileThreadResult: TileReturns unhandled."
compileThreadResult SegSpace
_ PatElem LetDecMem
_ RegTileReturns {} =
String -> ImpM MCMem HostEnv Multicore ()
forall a. String -> a
compilerBugS String
"compileThreadResult: RegTileReturns unhandled."
freeParams :: FreeIn a => a -> MulticoreGen [Imp.Param]
freeParams :: forall a. FreeIn a => a -> MulticoreGen [Param]
freeParams a
code = do
let free :: [VName]
free = Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ a -> Names
forall a. FreeIn a => a -> Names
freeIn a
code
[Type]
ts <- (VName -> ImpM MCMem HostEnv Multicore Type)
-> [VName] -> ImpM MCMem HostEnv Multicore [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> ImpM MCMem HostEnv Multicore Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType [VName]
free
[[Param]] -> [Param]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Param]] -> [Param])
-> ImpM MCMem HostEnv Multicore [[Param]] -> MulticoreGen [Param]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> Type -> MulticoreGen [Param])
-> [VName] -> [Type] -> ImpM MCMem HostEnv Multicore [[Param]]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM VName -> Type -> MulticoreGen [Param]
forall shape u. VName -> TypeBase shape u -> MulticoreGen [Param]
toParam [VName]
free [Type]
ts
groupResultArrays ::
String ->
SubExp ->
[SegBinOp MCMem] ->
MulticoreGen [[VName]]
groupResultArrays :: String -> SubExp -> [SegBinOp MCMem] -> MulticoreGen [[VName]]
groupResultArrays String
s SubExp
num_threads [SegBinOp MCMem]
reds =
[SegBinOp MCMem]
-> (SegBinOp MCMem -> ImpM MCMem HostEnv Multicore [VName])
-> MulticoreGen [[VName]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [SegBinOp MCMem]
reds ((SegBinOp MCMem -> ImpM MCMem HostEnv Multicore [VName])
-> MulticoreGen [[VName]])
-> (SegBinOp MCMem -> ImpM MCMem HostEnv Multicore [VName])
-> MulticoreGen [[VName]]
forall a b. (a -> b) -> a -> b
$ \(SegBinOp Commutativity
_ Lambda MCMem
lam [SubExp]
_ Shape
shape) ->
[Type]
-> (Type -> ImpM MCMem HostEnv Multicore VName)
-> ImpM MCMem HostEnv Multicore [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (Lambda MCMem -> [Type]
forall rep. Lambda rep -> [Type]
lambdaReturnType Lambda MCMem
lam) ((Type -> ImpM MCMem HostEnv Multicore VName)
-> ImpM MCMem HostEnv Multicore [VName])
-> (Type -> ImpM MCMem HostEnv Multicore VName)
-> ImpM MCMem HostEnv Multicore [VName]
forall a b. (a -> b) -> a -> b
$ \Type
t -> do
let full_shape :: Shape
full_shape = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
num_threads] Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Shape
shape Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape Type
t
String
-> PrimType -> Shape -> Space -> ImpM MCMem HostEnv Multicore VName
forall rep r op.
String -> PrimType -> Shape -> Space -> ImpM rep r op VName
sAllocArray String
s (Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType Type
t) Shape
full_shape Space
DefaultSpace
isLoadBalanced :: Imp.MCCode -> Bool
isLoadBalanced :: Code Multicore -> Bool
isLoadBalanced (Code Multicore
a Imp.:>>: Code Multicore
b) = Code Multicore -> Bool
isLoadBalanced Code Multicore
a Bool -> Bool -> Bool
&& Code Multicore -> Bool
isLoadBalanced Code Multicore
b
isLoadBalanced (Imp.For VName
_ Exp
_ Code Multicore
a) = Code Multicore -> Bool
isLoadBalanced Code Multicore
a
isLoadBalanced (Imp.If TExp Bool
_ Code Multicore
a Code Multicore
b) = Code Multicore -> Bool
isLoadBalanced Code Multicore
a Bool -> Bool -> Bool
&& Code Multicore -> Bool
isLoadBalanced Code Multicore
b
isLoadBalanced (Imp.Comment String
_ Code Multicore
a) = Code Multicore -> Bool
isLoadBalanced Code Multicore
a
isLoadBalanced Imp.While {} = Bool
False
isLoadBalanced (Imp.Op (Imp.ParLoop String
_ Code Multicore
code [Param]
_)) = Code Multicore -> Bool
isLoadBalanced Code Multicore
code
isLoadBalanced (Imp.Op (Imp.ForEachActive VName
_ Code Multicore
a)) = Code Multicore -> Bool
isLoadBalanced Code Multicore
a
isLoadBalanced (Imp.Op (Imp.ForEach VName
_ Exp
_ Exp
_ Code Multicore
a)) = Code Multicore -> Bool
isLoadBalanced Code Multicore
a
isLoadBalanced (Imp.Op (Imp.ISPCKernel Code Multicore
a [Param]
_)) = Code Multicore -> Bool
isLoadBalanced Code Multicore
a
isLoadBalanced Code Multicore
_ = Bool
True
segBinOpComm' :: [SegBinOp rep] -> Commutativity
segBinOpComm' :: forall rep. [SegBinOp rep] -> Commutativity
segBinOpComm' = [Commutativity] -> Commutativity
forall a. Monoid a => [a] -> a
mconcat ([Commutativity] -> Commutativity)
-> ([SegBinOp rep] -> [Commutativity])
-> [SegBinOp rep]
-> Commutativity
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SegBinOp rep -> Commutativity)
-> [SegBinOp rep] -> [Commutativity]
forall a b. (a -> b) -> [a] -> [b]
map SegBinOp rep -> Commutativity
forall rep. SegBinOp rep -> Commutativity
segBinOpComm
decideScheduling' :: SegOp () rep -> Imp.MCCode -> Imp.Scheduling
decideScheduling' :: forall rep. SegOp () rep -> Code Multicore -> Scheduling
decideScheduling' SegHist {} Code Multicore
_ = Scheduling
Imp.Static
decideScheduling' SegScan {} Code Multicore
_ = Scheduling
Imp.Static
decideScheduling' (SegRed ()
_ SegSpace
_ [SegBinOp rep]
reds [Type]
_ KernelBody rep
_) Code Multicore
code =
case [SegBinOp rep] -> Commutativity
forall rep. [SegBinOp rep] -> Commutativity
segBinOpComm' [SegBinOp rep]
reds of
Commutativity
Commutative -> Code Multicore -> Scheduling
decideScheduling Code Multicore
code
Commutativity
Noncommutative -> Scheduling
Imp.Static
decideScheduling' SegMap {} Code Multicore
code = Code Multicore -> Scheduling
decideScheduling Code Multicore
code
decideScheduling :: Imp.MCCode -> Imp.Scheduling
decideScheduling :: Code Multicore -> Scheduling
decideScheduling Code Multicore
code =
if Code Multicore -> Bool
isLoadBalanced Code Multicore
code
then Scheduling
Imp.Static
else Scheduling
Imp.Dynamic
extractAllocations :: Imp.MCCode -> (Imp.MCCode, Imp.MCCode)
Code Multicore
segop_code = Code Multicore -> (Code Multicore, Code Multicore)
forall {a}. Code Multicore -> (Code a, Code Multicore)
f Code Multicore
segop_code
where
declared :: Names
declared = Code Multicore -> Names
forall a. Code a -> Names
Imp.declaredIn Code Multicore
segop_code
f :: Code Multicore -> (Code a, Code Multicore)
f (Imp.DeclareMem VName
name Space
space) =
(VName -> Space -> Code a
forall a. VName -> Space -> Code a
Imp.DeclareMem VName
name Space
space, Code Multicore
forall a. Monoid a => a
mempty)
f (Imp.Allocate VName
name Count Bytes (TPrimExp Int64 VName)
size Space
space)
| Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Count Bytes (TPrimExp Int64 VName) -> Names
forall a. FreeIn a => a -> Names
freeIn Count Bytes (TPrimExp Int64 VName)
size Names -> Names -> Bool
`namesIntersect` Names
declared =
(VName -> Count Bytes (TPrimExp Int64 VName) -> Space -> Code a
forall a.
VName -> Count Bytes (TPrimExp Int64 VName) -> Space -> Code a
Imp.Allocate VName
name Count Bytes (TPrimExp Int64 VName)
size Space
space, Code Multicore
forall a. Monoid a => a
mempty)
f (Code Multicore
x Imp.:>>: Code Multicore
y) = Code Multicore -> (Code a, Code Multicore)
f Code Multicore
x (Code a, Code Multicore)
-> (Code a, Code Multicore) -> (Code a, Code Multicore)
forall a. Semigroup a => a -> a -> a
<> Code Multicore -> (Code a, Code Multicore)
f Code Multicore
y
f (Imp.While TExp Bool
cond Code Multicore
body) =
(Code a
forall a. Monoid a => a
mempty, TExp Bool -> Code Multicore -> Code Multicore
forall a. TExp Bool -> Code a -> Code a
Imp.While TExp Bool
cond Code Multicore
body)
f (Imp.For VName
i Exp
bound Code Multicore
body) =
(Code a
forall a. Monoid a => a
mempty, VName -> Exp -> Code Multicore -> Code Multicore
forall a. VName -> Exp -> Code a -> Code a
Imp.For VName
i Exp
bound Code Multicore
body)
f (Imp.Comment String
s Code Multicore
code) =
(Code Multicore -> Code Multicore)
-> (Code a, Code Multicore) -> (Code a, Code Multicore)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (String -> Code Multicore -> Code Multicore
forall a. String -> Code a -> Code a
Imp.Comment String
s) (Code Multicore -> (Code a, Code Multicore)
f Code Multicore
code)
f Imp.Free {} =
(Code a, Code Multicore)
forall a. Monoid a => a
mempty
f (Imp.If TExp Bool
cond Code Multicore
tcode Code Multicore
fcode) =
let (Code a
ta, Code Multicore
tcode') = Code Multicore -> (Code a, Code Multicore)
f Code Multicore
tcode
(Code a
fa, Code Multicore
fcode') = Code Multicore -> (Code a, Code Multicore)
f Code Multicore
fcode
in (Code a
ta Code a -> Code a -> Code a
forall a. Semigroup a => a -> a -> a
<> Code a
fa, TExp Bool -> Code Multicore -> Code Multicore -> Code Multicore
forall a. TExp Bool -> Code a -> Code a -> Code a
Imp.If TExp Bool
cond Code Multicore
tcode' Code Multicore
fcode')
f (Imp.Op (Imp.ParLoop String
s Code Multicore
body [Param]
free)) =
let (Code Multicore
body_allocs, Code Multicore
body') = Code Multicore -> (Code Multicore, Code Multicore)
extractAllocations Code Multicore
body
(Code a
free_allocs, Code Multicore
here_allocs) = Code Multicore -> (Code a, Code Multicore)
f Code Multicore
body_allocs
free' :: [Param]
free' =
(Param -> Bool) -> [Param] -> [Param]
forall a. (a -> Bool) -> [a] -> [a]
filter
( (VName -> Names -> Bool
`notNameIn` Code Multicore -> Names
forall a. Code a -> Names
Imp.declaredIn Code Multicore
body_allocs) (VName -> Bool) -> (Param -> VName) -> Param -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param -> VName
Imp.paramName
)
[Param]
free
in ( Code a
free_allocs,
Code Multicore
here_allocs Code Multicore -> Code Multicore -> Code Multicore
forall a. Semigroup a => a -> a -> a
<> Multicore -> Code Multicore
forall a. a -> Code a
Imp.Op (String -> Code Multicore -> [Param] -> Multicore
Imp.ParLoop String
s Code Multicore
body' [Param]
free')
)
f Code Multicore
code =
(Code a
forall a. Monoid a => a
mempty, Code Multicore
code)
data ChunkLoopVectorization = Vectorized | Scalar
generateChunkLoop ::
String ->
ChunkLoopVectorization ->
(Imp.TExp Int64 -> MulticoreGen ()) ->
MulticoreGen ()
generateChunkLoop :: String
-> ChunkLoopVectorization
-> (TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
generateChunkLoop String
desc ChunkLoopVectorization
Scalar TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ()
m = do
(TPrimExp Int64 VName
start, TPrimExp Int64 VName
end) <- MulticoreGen (TPrimExp Int64 VName, TPrimExp Int64 VName)
getLoopBounds
TPrimExp Int64 VName
n <- String
-> TPrimExp Int64 VName -> MulticoreGen (TPrimExp Int64 VName)
forall t rep r op. String -> TExp t -> ImpM rep r op (TExp t)
dPrimVE String
"n" (TPrimExp Int64 VName -> MulticoreGen (TPrimExp Int64 VName))
-> TPrimExp Int64 VName -> MulticoreGen (TPrimExp Int64 VName)
forall a b. (a -> b) -> a -> b
$ TPrimExp Int64 VName
end TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- TPrimExp Int64 VName
start
VName
i <- String -> ImpM MCMem HostEnv Multicore VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (String
desc String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"_i")
(Code Multicore
body_allocs, Code Multicore
body) <- (Code Multicore -> (Code Multicore, Code Multicore))
-> ImpM MCMem HostEnv Multicore (Code Multicore)
-> ImpM MCMem HostEnv Multicore (Code Multicore, Code Multicore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Code Multicore -> (Code Multicore, Code Multicore)
extractAllocations (ImpM MCMem HostEnv Multicore (Code Multicore)
-> ImpM MCMem HostEnv Multicore (Code Multicore, Code Multicore))
-> ImpM MCMem HostEnv Multicore (Code Multicore)
-> ImpM MCMem HostEnv Multicore (Code Multicore, Code Multicore)
forall a b. (a -> b) -> a -> b
$
ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore (Code Multicore)
forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
collect (ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore (Code Multicore))
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore (Code Multicore)
forall a b. (a -> b) -> a -> b
$ do
VName -> IntType -> ImpM MCMem HostEnv Multicore ()
forall rep r op. VName -> IntType -> ImpM rep r op ()
addLoopVar VName
i IntType
Int64
TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ()
m (TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ())
-> TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ TPrimExp Int64 VName
start TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 VName
i
Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall op rep r. Code op -> ImpM rep r op ()
emit Code Multicore
body_allocs
let bound :: Exp
bound = TPrimExp Int64 VName -> Exp
forall t v. TPrimExp t v -> PrimExp v
untyped TPrimExp Int64 VName
n
Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall op rep r. Code op -> ImpM rep r op ()
emit (Code Multicore -> ImpM MCMem HostEnv Multicore ())
-> Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ VName -> Exp -> Code Multicore -> Code Multicore
forall a. VName -> Exp -> Code a -> Code a
Imp.For VName
i Exp
bound Code Multicore
body
generateChunkLoop String
desc ChunkLoopVectorization
Vectorized TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ()
m = do
(TPrimExp Int64 VName
start, TPrimExp Int64 VName
end) <- MulticoreGen (TPrimExp Int64 VName, TPrimExp Int64 VName)
getLoopBounds
TPrimExp Int64 VName
n <- String
-> TPrimExp Int64 VName -> MulticoreGen (TPrimExp Int64 VName)
forall t rep r op. String -> TExp t -> ImpM rep r op (TExp t)
dPrimVE String
"n" (TPrimExp Int64 VName -> MulticoreGen (TPrimExp Int64 VName))
-> TPrimExp Int64 VName -> MulticoreGen (TPrimExp Int64 VName)
forall a b. (a -> b) -> a -> b
$ TPrimExp Int64 VName
end TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- TPrimExp Int64 VName
start
VName
i <- String -> ImpM MCMem HostEnv Multicore VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (String
desc String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"_i")
(Code Multicore
body_allocs, Code Multicore
body) <- (Code Multicore -> (Code Multicore, Code Multicore))
-> ImpM MCMem HostEnv Multicore (Code Multicore)
-> ImpM MCMem HostEnv Multicore (Code Multicore, Code Multicore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Code Multicore -> (Code Multicore, Code Multicore)
extractAllocations (ImpM MCMem HostEnv Multicore (Code Multicore)
-> ImpM MCMem HostEnv Multicore (Code Multicore, Code Multicore))
-> ImpM MCMem HostEnv Multicore (Code Multicore)
-> ImpM MCMem HostEnv Multicore (Code Multicore, Code Multicore)
forall a b. (a -> b) -> a -> b
$
ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore (Code Multicore)
forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
collect (ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore (Code Multicore))
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore (Code Multicore)
forall a b. (a -> b) -> a -> b
$ do
VName -> IntType -> ImpM MCMem HostEnv Multicore ()
forall rep r op. VName -> IntType -> ImpM rep r op ()
addLoopVar VName
i IntType
Int64
TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ()
m (TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ())
-> TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 VName
i
Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall op rep r. Code op -> ImpM rep r op ()
emit Code Multicore
body_allocs
let from :: Exp
from = TPrimExp Int64 VName -> Exp
forall t v. TPrimExp t v -> PrimExp v
untyped TPrimExp Int64 VName
start
let bound :: Exp
bound = TPrimExp Int64 VName -> Exp
forall t v. TPrimExp t v -> PrimExp v
untyped (TPrimExp Int64 VName
start TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int64 VName
n)
Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall op rep r. Code op -> ImpM rep r op ()
emit (Code Multicore -> ImpM MCMem HostEnv Multicore ())
-> Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ Multicore -> Code Multicore
forall a. a -> Code a
Imp.Op (Multicore -> Code Multicore) -> Multicore -> Code Multicore
forall a b. (a -> b) -> a -> b
$ VName -> Exp -> Exp -> Code Multicore -> Multicore
Imp.ForEach VName
i Exp
from Exp
bound Code Multicore
body
generateUniformizeLoop :: (Imp.TExp Int64 -> MulticoreGen ()) -> MulticoreGen ()
generateUniformizeLoop :: (TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
generateUniformizeLoop TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ()
m = do
VName
i <- String -> ImpM MCMem HostEnv Multicore VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"uni_i"
Code Multicore
body <- ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore (Code Multicore)
forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
collect (ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore (Code Multicore))
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore (Code Multicore)
forall a b. (a -> b) -> a -> b
$ do
VName -> IntType -> ImpM MCMem HostEnv Multicore ()
forall rep r op. VName -> IntType -> ImpM rep r op ()
addLoopVar VName
i IntType
Int64
TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ()
m (TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ())
-> TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 VName
i
Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall op rep r. Code op -> ImpM rep r op ()
emit (Code Multicore -> ImpM MCMem HostEnv Multicore ())
-> Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ Multicore -> Code Multicore
forall a. a -> Code a
Imp.Op (Multicore -> Code Multicore) -> Multicore -> Code Multicore
forall a b. (a -> b) -> a -> b
$ VName -> Code Multicore -> Multicore
Imp.ForEachActive VName
i Code Multicore
body
extractVectorLane :: Imp.TExp Int64 -> MulticoreGen Imp.MCCode -> MulticoreGen ()
TPrimExp Int64 VName
j ImpM MCMem HostEnv Multicore (Code Multicore)
code = do
let ut_exp :: Exp
ut_exp = TPrimExp Int64 VName -> Exp
forall t v. TPrimExp t v -> PrimExp v
untyped TPrimExp Int64 VName
j
Code Multicore
code' <- ImpM MCMem HostEnv Multicore (Code Multicore)
code
case Code Multicore
code' of
Imp.SetScalar VName
vname Exp
e -> do
Type
typ <- VName -> ImpM MCMem HostEnv Multicore Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
vname
case Type
typ of
Prim (FloatType FloatType
Float16) -> do
TV Any
tv <- String -> PrimType -> ImpM MCMem HostEnv Multicore (TV Any)
forall rep r op t. String -> PrimType -> ImpM rep r op (TV t)
dPrim String
"hack_extract_f16" (FloatType -> PrimType
FloatType FloatType
Float32)
Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall op rep r. Code op -> ImpM rep r op ()
emit (Code Multicore -> ImpM MCMem HostEnv Multicore ())
-> Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ VName -> Exp -> Code Multicore
forall a. VName -> Exp -> Code a
Imp.SetScalar (TV Any -> VName
forall t. TV t -> VName
tvVar TV Any
tv) Exp
e
Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall op rep r. Code op -> ImpM rep r op ()
emit (Code Multicore -> ImpM MCMem HostEnv Multicore ())
-> Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ Multicore -> Code Multicore
forall a. a -> Code a
Imp.Op (Multicore -> Code Multicore) -> Multicore -> Code Multicore
forall a b. (a -> b) -> a -> b
$ VName -> Exp -> Exp -> Multicore
Imp.ExtractLane VName
vname (TPrimExp Any VName -> Exp
forall t v. TPrimExp t v -> PrimExp v
untyped (TPrimExp Any VName -> Exp) -> TPrimExp Any VName -> Exp
forall a b. (a -> b) -> a -> b
$ TV Any -> TPrimExp Any VName
forall t. TV t -> TExp t
tvExp TV Any
tv) Exp
ut_exp
Type
_ -> Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall op rep r. Code op -> ImpM rep r op ()
emit (Code Multicore -> ImpM MCMem HostEnv Multicore ())
-> Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ Multicore -> Code Multicore
forall a. a -> Code a
Imp.Op (Multicore -> Code Multicore) -> Multicore -> Code Multicore
forall a b. (a -> b) -> a -> b
$ VName -> Exp -> Exp -> Multicore
Imp.ExtractLane VName
vname Exp
e Exp
ut_exp
Code Multicore
_ ->
Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall op rep r. Code op -> ImpM rep r op ()
emit Code Multicore
code'
inISPC :: MulticoreGen () -> MulticoreGen ()
inISPC :: ImpM MCMem HostEnv Multicore () -> ImpM MCMem HostEnv Multicore ()
inISPC ImpM MCMem HostEnv Multicore ()
code = do
Code Multicore
code' <- ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore (Code Multicore)
forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
collect ImpM MCMem HostEnv Multicore ()
code
[Param]
free <- Code Multicore -> MulticoreGen [Param]
forall a. FreeIn a => a -> MulticoreGen [Param]
freeParams Code Multicore
code'
Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall op rep r. Code op -> ImpM rep r op ()
emit (Code Multicore -> ImpM MCMem HostEnv Multicore ())
-> Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ Multicore -> Code Multicore
forall a. a -> Code a
Imp.Op (Multicore -> Code Multicore) -> Multicore -> Code Multicore
forall a b. (a -> b) -> a -> b
$ Code Multicore -> [Param] -> Multicore
Imp.ISPCKernel Code Multicore
code' [Param]
free
sForVectorized' :: VName -> Imp.Exp -> MulticoreGen () -> MulticoreGen ()
sForVectorized' :: VName
-> Exp
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ()
sForVectorized' VName
i Exp
bound ImpM MCMem HostEnv Multicore ()
body = do
let it :: IntType
it = case Exp -> PrimType
forall v. PrimExp v -> PrimType
primExpType Exp
bound of
IntType IntType
bound_t -> IntType
bound_t
PrimType
t -> String -> IntType
forall a. HasCallStack => String -> a
error (String -> IntType) -> String -> IntType
forall a b. (a -> b) -> a -> b
$ String
"sFor': bound " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Exp -> String
forall a. Pretty a => a -> String
pretty Exp
bound String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" is of type " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PrimType -> String
forall a. Pretty a => a -> String
pretty PrimType
t
VName -> IntType -> ImpM MCMem HostEnv Multicore ()
forall rep r op. VName -> IntType -> ImpM rep r op ()
addLoopVar VName
i IntType
it
Code Multicore
body' <- ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore (Code Multicore)
forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
collect ImpM MCMem HostEnv Multicore ()
body
Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall op rep r. Code op -> ImpM rep r op ()
emit (Code Multicore -> ImpM MCMem HostEnv Multicore ())
-> Code Multicore -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ Multicore -> Code Multicore
forall a. a -> Code a
Imp.Op (Multicore -> Code Multicore) -> Multicore -> Code Multicore
forall a b. (a -> b) -> a -> b
$ VName -> Exp -> Exp -> Code Multicore -> Multicore
Imp.ForEach VName
i (PrimValue -> Exp
forall v. PrimValue -> PrimExp v
Imp.ValueExp (PrimValue -> Exp) -> PrimValue -> Exp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimValue
blankPrimValue (PrimType -> PrimValue) -> PrimType -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Imp.IntType IntType
Imp.Int64) Exp
bound Code Multicore
body'
sForVectorized :: String -> Imp.TExp t -> (Imp.TExp t -> MulticoreGen ()) -> MulticoreGen ()
sForVectorized :: forall t.
String
-> TExp t
-> (TExp t -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
sForVectorized String
i TExp t
bound TExp t -> ImpM MCMem HostEnv Multicore ()
body = do
VName
i' <- String -> ImpM MCMem HostEnv Multicore VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
i
VName
-> Exp
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ()
sForVectorized' VName
i' (TExp t -> Exp
forall t v. TPrimExp t v -> PrimExp v
untyped TExp t
bound) (ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$
TExp t -> ImpM MCMem HostEnv Multicore ()
body (TExp t -> ImpM MCMem HostEnv Multicore ())
-> TExp t -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$
Exp -> TExp t
forall t v. PrimExp v -> TPrimExp t v
TPrimExp (Exp -> TExp t) -> Exp -> TExp t
forall a b. (a -> b) -> a -> b
$
VName -> PrimType -> Exp
Imp.var VName
i' (PrimType -> Exp) -> PrimType -> Exp
forall a b. (a -> b) -> a -> b
$
Exp -> PrimType
forall v. PrimExp v -> PrimType
primExpType (Exp -> PrimType) -> Exp -> PrimType
forall a b. (a -> b) -> a -> b
$
TExp t -> Exp
forall t v. TPrimExp t v -> PrimExp v
untyped TExp t
bound
sLoopNestVectorized ::
Shape ->
([Imp.TExp Int64] -> MulticoreGen ()) ->
MulticoreGen ()
sLoopNestVectorized :: Shape
-> ([TPrimExp Int64 VName] -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
sLoopNestVectorized = [TPrimExp Int64 VName]
-> [SubExp]
-> ([TPrimExp Int64 VName] -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
sLoopNest' [] ([SubExp]
-> ([TPrimExp Int64 VName] -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ())
-> (Shape -> [SubExp])
-> Shape
-> ([TPrimExp Int64 VName] -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims
where
sLoopNest' :: [TPrimExp Int64 VName]
-> [SubExp]
-> ([TPrimExp Int64 VName] -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
sLoopNest' [TPrimExp Int64 VName]
is [] [TPrimExp Int64 VName] -> ImpM MCMem HostEnv Multicore ()
f = [TPrimExp Int64 VName] -> ImpM MCMem HostEnv Multicore ()
f ([TPrimExp Int64 VName] -> ImpM MCMem HostEnv Multicore ())
-> [TPrimExp Int64 VName] -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. [a] -> [a]
reverse [TPrimExp Int64 VName]
is
sLoopNest' [TPrimExp Int64 VName]
is [SubExp
d] [TPrimExp Int64 VName] -> ImpM MCMem HostEnv Multicore ()
f =
String
-> TPrimExp Int64 VName
-> (TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
forall t.
String
-> TExp t
-> (TExp t -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
sForVectorized String
"nest_i" (SubExp -> TPrimExp Int64 VName
pe64 SubExp
d) ((TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ())
-> (TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ \TPrimExp Int64 VName
i -> [TPrimExp Int64 VName]
-> [SubExp]
-> ([TPrimExp Int64 VName] -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
sLoopNest' (TPrimExp Int64 VName
i TPrimExp Int64 VName
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. a -> [a] -> [a]
: [TPrimExp Int64 VName]
is) [] [TPrimExp Int64 VName] -> ImpM MCMem HostEnv Multicore ()
f
sLoopNest' [TPrimExp Int64 VName]
is (SubExp
d : [SubExp]
ds) [TPrimExp Int64 VName] -> ImpM MCMem HostEnv Multicore ()
f =
String
-> TPrimExp Int64 VName
-> (TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
forall t rep r op.
String
-> TExp t -> (TExp t -> ImpM rep r op ()) -> ImpM rep r op ()
sFor String
"nest_i" (SubExp -> TPrimExp Int64 VName
pe64 SubExp
d) ((TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ())
-> (TPrimExp Int64 VName -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ \TPrimExp Int64 VName
i -> [TPrimExp Int64 VName]
-> [SubExp]
-> ([TPrimExp Int64 VName] -> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
sLoopNest' (TPrimExp Int64 VName
i TPrimExp Int64 VName
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. a -> [a] -> [a]
: [TPrimExp Int64 VName]
is) [SubExp]
ds [TPrimExp Int64 VName] -> ImpM MCMem HostEnv Multicore ()
f
renameHistOpLambda :: [HistOp MCMem] -> MulticoreGen [HistOp MCMem]
renameHistOpLambda :: [HistOp MCMem] -> MulticoreGen [HistOp MCMem]
renameHistOpLambda [HistOp MCMem]
hist_ops =
[HistOp MCMem]
-> (HistOp MCMem -> ImpM MCMem HostEnv Multicore (HistOp MCMem))
-> MulticoreGen [HistOp MCMem]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [HistOp MCMem]
hist_ops ((HistOp MCMem -> ImpM MCMem HostEnv Multicore (HistOp MCMem))
-> MulticoreGen [HistOp MCMem])
-> (HistOp MCMem -> ImpM MCMem HostEnv Multicore (HistOp MCMem))
-> MulticoreGen [HistOp MCMem]
forall a b. (a -> b) -> a -> b
$ \(HistOp Shape
w SubExp
rf [VName]
dest [SubExp]
neutral Shape
shape Lambda MCMem
lam) -> do
Lambda MCMem
lam' <- Lambda MCMem -> ImpM MCMem HostEnv Multicore (Lambda MCMem)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Lambda rep -> m (Lambda rep)
renameLambda Lambda MCMem
lam
HistOp MCMem -> ImpM MCMem HostEnv Multicore (HistOp MCMem)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (HistOp MCMem -> ImpM MCMem HostEnv Multicore (HistOp MCMem))
-> HistOp MCMem -> ImpM MCMem HostEnv Multicore (HistOp MCMem)
forall a b. (a -> b) -> a -> b
$ Shape
-> SubExp
-> [VName]
-> [SubExp]
-> Shape
-> Lambda MCMem
-> HistOp MCMem
forall rep.
Shape
-> SubExp
-> [VName]
-> [SubExp]
-> Shape
-> Lambda rep
-> HistOp rep
HistOp Shape
w SubExp
rf [VName]
dest [SubExp]
neutral Shape
shape Lambda MCMem
lam'
data Locking = Locking
{
Locking -> VName
lockingArray :: VName,
Locking -> TExp Int32
lockingIsUnlocked :: Imp.TExp Int32,
Locking -> TExp Int32
lockingToLock :: Imp.TExp Int32,
Locking -> TExp Int32
lockingToUnlock :: Imp.TExp Int32,
Locking -> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
lockingMapping :: [Imp.TExp Int64] -> [Imp.TExp Int64]
}
type DoAtomicUpdate rep r =
[VName] -> [Imp.TExp Int64] -> MulticoreGen ()
data AtomicUpdate rep r
= AtomicPrim (DoAtomicUpdate rep r)
|
AtomicCAS (DoAtomicUpdate rep r)
|
AtomicLocking (Locking -> DoAtomicUpdate rep r)
atomicUpdateLocking ::
AtomicBinOp ->
Lambda MCMem ->
AtomicUpdate MCMem ()
atomicUpdateLocking :: AtomicBinOp -> Lambda MCMem -> AtomicUpdate MCMem ()
atomicUpdateLocking AtomicBinOp
atomicBinOp Lambda MCMem
lam
| Just [(BinOp, PrimType, VName, VName)]
ops_and_ts <- Lambda MCMem -> Maybe [(BinOp, PrimType, VName, VName)]
forall rep.
ASTRep rep =>
Lambda rep -> Maybe [(BinOp, PrimType, VName, VName)]
lamIsBinOp Lambda MCMem
lam,
((BinOp, PrimType, VName, VName) -> Bool)
-> [(BinOp, PrimType, VName, VName)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(BinOp
_, PrimType
t, VName
_, VName
_) -> Int -> Bool
supportedPrims (Int -> Bool) -> Int -> Bool
forall a b. (a -> b) -> a -> b
$ PrimType -> Int
primBitSize PrimType
t) [(BinOp, PrimType, VName, VName)]
ops_and_ts =
[(BinOp, PrimType, VName, VName)]
-> DoAtomicUpdate MCMem () -> AtomicUpdate MCMem ()
forall {t :: * -> *} {b} {c} {d} {rep} {r}.
Foldable t =>
t (BinOp, b, c, d) -> DoAtomicUpdate MCMem () -> AtomicUpdate rep r
primOrCas [(BinOp, PrimType, VName, VName)]
ops_and_ts (DoAtomicUpdate MCMem () -> AtomicUpdate MCMem ())
-> DoAtomicUpdate MCMem () -> AtomicUpdate MCMem ()
forall a b. (a -> b) -> a -> b
$ \[VName]
arrs [TPrimExp Int64 VName]
bucket ->
[(VName, (BinOp, PrimType, VName, VName))]
-> ((VName, (BinOp, PrimType, VName, VName))
-> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName]
-> [(BinOp, PrimType, VName, VName)]
-> [(VName, (BinOp, PrimType, VName, VName))]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
arrs [(BinOp, PrimType, VName, VName)]
ops_and_ts) (((VName, (BinOp, PrimType, VName, VName))
-> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ())
-> ((VName, (BinOp, PrimType, VName, VName))
-> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ \(VName
a, (BinOp
op, PrimType
t, VName
x, VName
y)) -> do
TV Any
old <- String -> PrimType -> ImpM MCMem HostEnv Multicore (TV Any)
forall rep r op t. String -> PrimType -> ImpM rep r op (TV t)
dPrim String
"old" PrimType
t
(VName
arr', Space
_a_space, Count Elements (TPrimExp Int64 VName)
bucket_offset) <- VName
-> [TPrimExp Int64 VName]
-> ImpM
MCMem
HostEnv
Multicore
(VName, Space, Count Elements (TPrimExp Int64 VName))
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> ImpM
rep r op (VName, Space, Count Elements (TPrimExp Int64 VName))
fullyIndexArray VName
a [TPrimExp Int64 VName]
bucket
case VName
-> VName
-> Count Elements (TExp Int32)
-> BinOp
-> Maybe (Exp -> Multicore)
opHasAtomicSupport (TV Any -> VName
forall t. TV t -> VName
tvVar TV Any
old) VName
arr' (TPrimExp Int64 VName -> TExp Int32
forall t v. IntExp t => TPrimExp t v -> TPrimExp Int32 v
sExt32 (TPrimExp Int64 VName -> TExp Int32)
-> Count Elements (TPrimExp Int64 VName)
-> Count Elements (TExp Int32)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Count Elements (TPrimExp Int64 VName)
bucket_offset) BinOp
op of
Just Exp -> Multicore
f -> Multicore -> ImpM MCMem HostEnv Multicore ()
forall op rep r. op -> ImpM rep r op ()
sOp (Multicore -> ImpM MCMem HostEnv Multicore ())
-> Multicore -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ Exp -> Multicore
f (Exp -> Multicore) -> Exp -> Multicore
forall a b. (a -> b) -> a -> b
$ VName -> PrimType -> Exp
Imp.var VName
y PrimType
t
Maybe (Exp -> Multicore)
Nothing ->
PrimType
-> VName
-> VName
-> [TPrimExp Int64 VName]
-> VName
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ()
atomicUpdateCAS PrimType
t VName
a (TV Any -> VName
forall t. TV t -> VName
tvVar TV Any
old) [TPrimExp Int64 VName]
bucket VName
x (ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$
VName
x VName -> Exp -> ImpM MCMem HostEnv Multicore ()
forall rep r op. VName -> Exp -> ImpM rep r op ()
<~~ BinOp -> Exp -> Exp -> Exp
forall v. BinOp -> PrimExp v -> PrimExp v -> PrimExp v
Imp.BinOpExp BinOp
op (VName -> PrimType -> Exp
Imp.var VName
x PrimType
t) (VName -> PrimType -> Exp
Imp.var VName
y PrimType
t)
where
opHasAtomicSupport :: VName
-> VName
-> Count Elements (TExp Int32)
-> BinOp
-> Maybe (Exp -> Multicore)
opHasAtomicSupport VName
old VName
arr' Count Elements (TExp Int32)
bucket' BinOp
bop = do
let atomic :: (VName -> VName -> Count Elements (TExp Int32) -> a -> AtomicOp)
-> a -> Multicore
atomic VName -> VName -> Count Elements (TExp Int32) -> a -> AtomicOp
f = AtomicOp -> Multicore
Imp.Atomic (AtomicOp -> Multicore) -> (a -> AtomicOp) -> a -> Multicore
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> VName -> Count Elements (TExp Int32) -> a -> AtomicOp
f VName
old VName
arr' Count Elements (TExp Int32)
bucket'
(VName -> VName -> Count Elements (TExp Int32) -> Exp -> AtomicOp)
-> Exp -> Multicore
forall {a}.
(VName -> VName -> Count Elements (TExp Int32) -> a -> AtomicOp)
-> a -> Multicore
atomic ((VName -> VName -> Count Elements (TExp Int32) -> Exp -> AtomicOp)
-> Exp -> Multicore)
-> Maybe
(VName -> VName -> Count Elements (TExp Int32) -> Exp -> AtomicOp)
-> Maybe (Exp -> Multicore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> AtomicBinOp
atomicBinOp BinOp
bop
primOrCas :: t (BinOp, b, c, d) -> DoAtomicUpdate MCMem () -> AtomicUpdate rep r
primOrCas t (BinOp, b, c, d)
ops
| ((BinOp, b, c, d) -> Bool) -> t (BinOp, b, c, d) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (BinOp, b, c, d) -> Bool
forall {b} {c} {d}. (BinOp, b, c, d) -> Bool
isPrim t (BinOp, b, c, d)
ops = DoAtomicUpdate MCMem () -> AtomicUpdate rep r
forall rep r. DoAtomicUpdate MCMem () -> AtomicUpdate rep r
AtomicPrim
| Bool
otherwise = DoAtomicUpdate MCMem () -> AtomicUpdate rep r
forall rep r. DoAtomicUpdate MCMem () -> AtomicUpdate rep r
AtomicCAS
isPrim :: (BinOp, b, c, d) -> Bool
isPrim (BinOp
op, b
_, c
_, d
_) = Maybe
(VName -> VName -> Count Elements (TExp Int32) -> Exp -> AtomicOp)
-> Bool
forall a. Maybe a -> Bool
isJust (Maybe
(VName -> VName -> Count Elements (TExp Int32) -> Exp -> AtomicOp)
-> Bool)
-> Maybe
(VName -> VName -> Count Elements (TExp Int32) -> Exp -> AtomicOp)
-> Bool
forall a b. (a -> b) -> a -> b
$ AtomicBinOp
atomicBinOp BinOp
op
atomicUpdateLocking AtomicBinOp
_ Lambda MCMem
op
| [Prim PrimType
t] <- Lambda MCMem -> [Type]
forall rep. Lambda rep -> [Type]
lambdaReturnType Lambda MCMem
op,
[LParam MCMem
xp, LParam MCMem
_] <- Lambda MCMem -> [LParam MCMem]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda MCMem
op,
Int -> Bool
supportedPrims (PrimType -> Int
primBitSize PrimType
t) = DoAtomicUpdate MCMem () -> AtomicUpdate MCMem ()
forall rep r. DoAtomicUpdate MCMem () -> AtomicUpdate rep r
AtomicCAS (DoAtomicUpdate MCMem () -> AtomicUpdate MCMem ())
-> DoAtomicUpdate MCMem () -> AtomicUpdate MCMem ()
forall a b. (a -> b) -> a -> b
$ \[VName
arr] [TPrimExp Int64 VName]
bucket -> do
TV Any
old <- String -> PrimType -> ImpM MCMem HostEnv Multicore (TV Any)
forall rep r op t. String -> PrimType -> ImpM rep r op (TV t)
dPrim String
"old" PrimType
t
PrimType
-> VName
-> VName
-> [TPrimExp Int64 VName]
-> VName
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ()
atomicUpdateCAS PrimType
t VName
arr (TV Any -> VName
forall t. TV t -> VName
tvVar TV Any
old) [TPrimExp Int64 VName]
bucket (Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName LParam MCMem
Param LetDecMem
xp) (ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$
[Param LetDecMem] -> Body MCMem -> ImpM MCMem HostEnv Multicore ()
forall dec rep r op. [Param dec] -> Body rep -> ImpM rep r op ()
compileBody' [LParam MCMem
Param LetDecMem
xp] (Body MCMem -> ImpM MCMem HostEnv Multicore ())
-> Body MCMem -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$
Lambda MCMem -> Body MCMem
forall rep. Lambda rep -> Body rep
lambdaBody Lambda MCMem
op
atomicUpdateLocking AtomicBinOp
_ Lambda MCMem
op = (Locking -> DoAtomicUpdate MCMem ()) -> AtomicUpdate MCMem ()
forall rep r.
(Locking -> DoAtomicUpdate MCMem ()) -> AtomicUpdate rep r
AtomicLocking ((Locking -> DoAtomicUpdate MCMem ()) -> AtomicUpdate MCMem ())
-> (Locking -> DoAtomicUpdate MCMem ()) -> AtomicUpdate MCMem ()
forall a b. (a -> b) -> a -> b
$ \Locking
locking [VName]
arrs [TPrimExp Int64 VName]
bucket -> do
TV Int32
old <- String -> PrimType -> ImpM MCMem HostEnv Multicore (TV Int32)
forall rep r op t. String -> PrimType -> ImpM rep r op (TV t)
dPrim String
"old" PrimType
int32
TV Int32
continue <- String
-> PrimType
-> TExp Int32
-> ImpM MCMem HostEnv Multicore (TV Int32)
forall t rep r op.
String -> PrimType -> TExp t -> ImpM rep r op (TV t)
dPrimVol String
"continue" PrimType
int32 (TExp Int32
0 :: Imp.TExp Int32)
(VName
locks', Space
_locks_space, Count Elements (TPrimExp Int64 VName)
locks_offset) <-
VName
-> [TPrimExp Int64 VName]
-> ImpM
MCMem
HostEnv
Multicore
(VName, Space, Count Elements (TPrimExp Int64 VName))
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> ImpM
rep r op (VName, Space, Count Elements (TPrimExp Int64 VName))
fullyIndexArray (Locking -> VName
lockingArray Locking
locking) ([TPrimExp Int64 VName]
-> ImpM
MCMem
HostEnv
Multicore
(VName, Space, Count Elements (TPrimExp Int64 VName)))
-> [TPrimExp Int64 VName]
-> ImpM
MCMem
HostEnv
Multicore
(VName, Space, Count Elements (TPrimExp Int64 VName))
forall a b. (a -> b) -> a -> b
$ Locking -> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
lockingMapping Locking
locking [TPrimExp Int64 VName]
bucket
let try_acquire_lock :: ImpM rep r Multicore ()
try_acquire_lock = do
TV Int32
old TV Int32 -> TExp Int32 -> ImpM rep r Multicore ()
forall t rep r op. TV t -> TExp t -> ImpM rep r op ()
<-- (TExp Int32
0 :: Imp.TExp Int32)
Multicore -> ImpM rep r Multicore ()
forall op rep r. op -> ImpM rep r op ()
sOp (Multicore -> ImpM rep r Multicore ())
-> (AtomicOp -> Multicore) -> AtomicOp -> ImpM rep r Multicore ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AtomicOp -> Multicore
Imp.Atomic (AtomicOp -> ImpM rep r Multicore ())
-> AtomicOp -> ImpM rep r Multicore ()
forall a b. (a -> b) -> a -> b
$
PrimType
-> VName
-> VName
-> Count Elements (TExp Int32)
-> VName
-> Exp
-> AtomicOp
Imp.AtomicCmpXchg
PrimType
int32
(TV Int32 -> VName
forall t. TV t -> VName
tvVar TV Int32
old)
VName
locks'
(TPrimExp Int64 VName -> TExp Int32
forall t v. IntExp t => TPrimExp t v -> TPrimExp Int32 v
sExt32 (TPrimExp Int64 VName -> TExp Int32)
-> Count Elements (TPrimExp Int64 VName)
-> Count Elements (TExp Int32)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Count Elements (TPrimExp Int64 VName)
locks_offset)
(TV Int32 -> VName
forall t. TV t -> VName
tvVar TV Int32
continue)
(TExp Int32 -> Exp
forall t v. TPrimExp t v -> PrimExp v
untyped (Locking -> TExp Int32
lockingToLock Locking
locking))
lock_acquired :: TExp Int32
lock_acquired = TV Int32 -> TExp Int32
forall t. TV t -> TExp t
tvExp TV Int32
continue
release_lock :: ImpM rep r Multicore ()
release_lock = do
TV Int32
old TV Int32 -> TExp Int32 -> ImpM rep r Multicore ()
forall t rep r op. TV t -> TExp t -> ImpM rep r op ()
<-- Locking -> TExp Int32
lockingToLock Locking
locking
Multicore -> ImpM rep r Multicore ()
forall op rep r. op -> ImpM rep r op ()
sOp (Multicore -> ImpM rep r Multicore ())
-> (AtomicOp -> Multicore) -> AtomicOp -> ImpM rep r Multicore ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AtomicOp -> Multicore
Imp.Atomic (AtomicOp -> ImpM rep r Multicore ())
-> AtomicOp -> ImpM rep r Multicore ()
forall a b. (a -> b) -> a -> b
$
PrimType
-> VName
-> VName
-> Count Elements (TExp Int32)
-> VName
-> Exp
-> AtomicOp
Imp.AtomicCmpXchg
PrimType
int32
(TV Int32 -> VName
forall t. TV t -> VName
tvVar TV Int32
old)
VName
locks'
(TPrimExp Int64 VName -> TExp Int32
forall t v. IntExp t => TPrimExp t v -> TPrimExp Int32 v
sExt32 (TPrimExp Int64 VName -> TExp Int32)
-> Count Elements (TPrimExp Int64 VName)
-> Count Elements (TExp Int32)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Count Elements (TPrimExp Int64 VName)
locks_offset)
(TV Int32 -> VName
forall t. TV t -> VName
tvVar TV Int32
continue)
(TExp Int32 -> Exp
forall t v. TPrimExp t v -> PrimExp v
untyped (Locking -> TExp Int32
lockingToUnlock Locking
locking))
let ([Param LetDecMem]
acc_params, [Param LetDecMem]
_arr_params) = Int -> [Param LetDecMem] -> ([Param LetDecMem], [Param LetDecMem])
forall a. Int -> [a] -> ([a], [a])
splitAt ([VName] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
arrs) ([Param LetDecMem] -> ([Param LetDecMem], [Param LetDecMem]))
-> [Param LetDecMem] -> ([Param LetDecMem], [Param LetDecMem])
forall a b. (a -> b) -> a -> b
$ Lambda MCMem -> [LParam MCMem]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda MCMem
op
bind_acc_params :: ImpM rep r op ()
bind_acc_params =
ImpM rep r op () -> ImpM rep r op ()
forall rep r op a. ImpM rep r op a -> ImpM rep r op a
everythingVolatile (ImpM rep r op () -> ImpM rep r op ())
-> ImpM rep r op () -> ImpM rep r op ()
forall a b. (a -> b) -> a -> b
$
String -> ImpM rep r op () -> ImpM rep r op ()
forall rep r op. String -> ImpM rep r op () -> ImpM rep r op ()
sComment String
"bind lhs" (ImpM rep r op () -> ImpM rep r op ())
-> ImpM rep r op () -> ImpM rep r op ()
forall a b. (a -> b) -> a -> b
$
[(Param LetDecMem, VName)]
-> ((Param LetDecMem, VName) -> ImpM rep r op ())
-> ImpM rep r op ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LetDecMem] -> [VName] -> [(Param LetDecMem, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param LetDecMem]
acc_params [VName]
arrs) (((Param LetDecMem, VName) -> ImpM rep r op ())
-> ImpM rep r op ())
-> ((Param LetDecMem, VName) -> ImpM rep r op ())
-> ImpM rep r op ()
forall a b. (a -> b) -> a -> b
$ \(Param LetDecMem
acc_p, VName
arr) ->
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName Param LetDecMem
acc_p) [] (VName -> SubExp
Var VName
arr) [TPrimExp Int64 VName]
bucket
let op_body :: ImpM MCMem r op ()
op_body =
String -> ImpM MCMem r op () -> ImpM MCMem r op ()
forall rep r op. String -> ImpM rep r op () -> ImpM rep r op ()
sComment String
"execute operation" (ImpM MCMem r op () -> ImpM MCMem r op ())
-> ImpM MCMem r op () -> ImpM MCMem r op ()
forall a b. (a -> b) -> a -> b
$
[Param LetDecMem] -> Body MCMem -> ImpM MCMem r op ()
forall dec rep r op. [Param dec] -> Body rep -> ImpM rep r op ()
compileBody' [Param LetDecMem]
acc_params (Body MCMem -> ImpM MCMem r op ())
-> Body MCMem -> ImpM MCMem r op ()
forall a b. (a -> b) -> a -> b
$
Lambda MCMem -> Body MCMem
forall rep. Lambda rep -> Body rep
lambdaBody Lambda MCMem
op
do_hist :: ImpM rep r op ()
do_hist =
ImpM rep r op () -> ImpM rep r op ()
forall rep r op a. ImpM rep r op a -> ImpM rep r op a
everythingVolatile (ImpM rep r op () -> ImpM rep r op ())
-> ImpM rep r op () -> ImpM rep r op ()
forall a b. (a -> b) -> a -> b
$
String -> ImpM rep r op () -> ImpM rep r op ()
forall rep r op. String -> ImpM rep r op () -> ImpM rep r op ()
sComment String
"update global result" (ImpM rep r op () -> ImpM rep r op ())
-> ImpM rep r op () -> ImpM rep r op ()
forall a b. (a -> b) -> a -> b
$
(VName -> SubExp -> ImpM rep r op ())
-> [VName] -> [SubExp] -> ImpM rep r op ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ ([TPrimExp Int64 VName] -> VName -> SubExp -> ImpM rep r op ()
forall {rep} {r} {op}.
[TPrimExp Int64 VName] -> VName -> SubExp -> ImpM rep r op ()
writeArray [TPrimExp Int64 VName]
bucket) [VName]
arrs ([SubExp] -> ImpM rep r op ()) -> [SubExp] -> ImpM rep r op ()
forall a b. (a -> b) -> a -> b
$
(Param LetDecMem -> SubExp) -> [Param LetDecMem] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
Var (VName -> SubExp)
-> (Param LetDecMem -> VName) -> Param LetDecMem -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName) [Param LetDecMem]
acc_params
TExp Bool
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ()
forall rep r op. TExp Bool -> ImpM rep r op () -> ImpM rep r op ()
sWhile (TV Int32 -> TExp Int32
forall t. TV t -> TExp t
tvExp TV Int32
continue TExp Int32 -> TExp Int32 -> TExp Bool
forall t v. TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. TExp Int32
0) (ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ do
ImpM MCMem HostEnv Multicore ()
forall {rep} {r}. ImpM rep r Multicore ()
try_acquire_lock
TExp Bool
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ()
forall rep r op. TExp Bool -> ImpM rep r op () -> ImpM rep r op ()
sUnless (TExp Int32
lock_acquired TExp Int32 -> TExp Int32 -> TExp Bool
forall t v. TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. TExp Int32
0) (ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ do
[LParam MCMem] -> ImpM MCMem HostEnv Multicore ()
forall rep inner r op.
Mem rep inner =>
[LParam rep] -> ImpM rep r op ()
dLParams [LParam MCMem]
[Param LetDecMem]
acc_params
ImpM MCMem HostEnv Multicore ()
forall {rep} {r} {op}. ImpM rep r op ()
bind_acc_params
ImpM MCMem HostEnv Multicore ()
forall {r} {op}. ImpM MCMem r op ()
op_body
ImpM MCMem HostEnv Multicore ()
forall {rep} {r} {op}. ImpM rep r op ()
do_hist
ImpM MCMem HostEnv Multicore ()
forall {rep} {r}. ImpM rep r Multicore ()
release_lock
where
writeArray :: [TPrimExp Int64 VName] -> VName -> SubExp -> ImpM rep r op ()
writeArray [TPrimExp Int64 VName]
bucket VName
arr SubExp
val = VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix VName
arr [TPrimExp Int64 VName]
bucket SubExp
val []
atomicUpdateCAS ::
PrimType ->
VName ->
VName ->
[Imp.TExp Int64] ->
VName ->
MulticoreGen () ->
MulticoreGen ()
atomicUpdateCAS :: PrimType
-> VName
-> VName
-> [TPrimExp Int64 VName]
-> VName
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ()
atomicUpdateCAS PrimType
t VName
arr VName
old [TPrimExp Int64 VName]
bucket VName
x ImpM MCMem HostEnv Multicore ()
do_op = do
TV Int32
run_loop <- String -> TExp Int32 -> ImpM MCMem HostEnv Multicore (TV Int32)
forall t rep r op. String -> TExp t -> ImpM rep r op (TV t)
dPrimV String
"run_loop" (TExp Int32
0 :: Imp.TExp Int32)
(VName
arr', Space
_a_space, Count Elements (TPrimExp Int64 VName)
bucket_offset) <- VName
-> [TPrimExp Int64 VName]
-> ImpM
MCMem
HostEnv
Multicore
(VName, Space, Count Elements (TPrimExp Int64 VName))
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> ImpM
rep r op (VName, Space, Count Elements (TPrimExp Int64 VName))
fullyIndexArray VName
arr [TPrimExp Int64 VName]
bucket
PrimType
bytes <- Int -> MulticoreGen PrimType
toIntegral (Int -> MulticoreGen PrimType) -> Int -> MulticoreGen PrimType
forall a b. (a -> b) -> a -> b
$ PrimType -> Int
primBitSize PrimType
t
let (PrimExp v -> PrimExp v
toBits, PrimExp v -> PrimExp v
fromBits) =
case PrimType
t of
FloatType FloatType
Float16 ->
( \PrimExp v
v -> String -> [PrimExp v] -> PrimType -> PrimExp v
forall v. String -> [PrimExp v] -> PrimType -> PrimExp v
Imp.FunExp String
"to_bits16" [PrimExp v
v] PrimType
int16,
\PrimExp v
v -> String -> [PrimExp v] -> PrimType -> PrimExp v
forall v. String -> [PrimExp v] -> PrimType -> PrimExp v
Imp.FunExp String
"from_bits16" [PrimExp v
v] PrimType
t
)
FloatType FloatType
Float32 ->
( \PrimExp v
v -> String -> [PrimExp v] -> PrimType -> PrimExp v
forall v. String -> [PrimExp v] -> PrimType -> PrimExp v
Imp.FunExp String
"to_bits32" [PrimExp v
v] PrimType
int32,
\PrimExp v
v -> String -> [PrimExp v] -> PrimType -> PrimExp v
forall v. String -> [PrimExp v] -> PrimType -> PrimExp v
Imp.FunExp String
"from_bits32" [PrimExp v
v] PrimType
t
)
FloatType FloatType
Float64 ->
( \PrimExp v
v -> String -> [PrimExp v] -> PrimType -> PrimExp v
forall v. String -> [PrimExp v] -> PrimType -> PrimExp v
Imp.FunExp String
"to_bits64" [PrimExp v
v] PrimType
int64,
\PrimExp v
v -> String -> [PrimExp v] -> PrimType -> PrimExp v
forall v. String -> [PrimExp v] -> PrimType -> PrimExp v
Imp.FunExp String
"from_bits64" [PrimExp v
v] PrimType
t
)
PrimType
_ -> (PrimExp v -> PrimExp v
forall a. a -> a
id, PrimExp v -> PrimExp v
forall a. a -> a
id)
int :: PrimType
int
| PrimType -> Int
primBitSize PrimType
t Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
16 = PrimType
int16
| PrimType -> Int
primBitSize PrimType
t Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
32 = PrimType
int32
| Bool
otherwise = PrimType
int64
ImpM MCMem HostEnv Multicore () -> ImpM MCMem HostEnv Multicore ()
forall rep r op a. ImpM rep r op a -> ImpM rep r op a
everythingVolatile (ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM MCMem HostEnv Multicore ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix VName
old [] (VName -> SubExp
Var VName
arr) [TPrimExp Int64 VName]
bucket
VName
old_bits_v <- TV Any -> VName
forall t. TV t -> VName
tvVar (TV Any -> VName)
-> ImpM MCMem HostEnv Multicore (TV Any)
-> ImpM MCMem HostEnv Multicore VName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> PrimType -> ImpM MCMem HostEnv Multicore (TV Any)
forall rep r op t. String -> PrimType -> ImpM rep r op (TV t)
dPrim String
"old_bits" PrimType
int
VName
old_bits_v VName -> Exp -> ImpM MCMem HostEnv Multicore ()
forall rep r op. VName -> Exp -> ImpM rep r op ()
<~~ Exp -> Exp
forall {v}. PrimExp v -> PrimExp v
toBits (VName -> PrimType -> Exp
Imp.var VName
old PrimType
t)
let old_bits :: Exp
old_bits = VName -> PrimType -> Exp
Imp.var VName
old_bits_v PrimType
int
TExp Bool
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ()
forall rep r op. TExp Bool -> ImpM rep r op () -> ImpM rep r op ()
sWhile (TV Int32 -> TExp Int32
forall t. TV t -> TExp t
tvExp TV Int32
run_loop TExp Int32 -> TExp Int32 -> TExp Bool
forall t v. TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. TExp Int32
0) (ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ())
-> ImpM MCMem HostEnv Multicore ()
-> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$ do
VName
x VName -> Exp -> ImpM MCMem HostEnv Multicore ()
forall rep r op. VName -> Exp -> ImpM rep r op ()
<~~ VName -> PrimType -> Exp
Imp.var VName
old PrimType
t
ImpM MCMem HostEnv Multicore ()
do_op
Multicore -> ImpM MCMem HostEnv Multicore ()
forall op rep r. op -> ImpM rep r op ()
sOp (Multicore -> ImpM MCMem HostEnv Multicore ())
-> (AtomicOp -> Multicore)
-> AtomicOp
-> ImpM MCMem HostEnv Multicore ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AtomicOp -> Multicore
Imp.Atomic (AtomicOp -> ImpM MCMem HostEnv Multicore ())
-> AtomicOp -> ImpM MCMem HostEnv Multicore ()
forall a b. (a -> b) -> a -> b
$
PrimType
-> VName
-> VName
-> Count Elements (TExp Int32)
-> VName
-> Exp
-> AtomicOp
Imp.AtomicCmpXchg
PrimType
bytes
VName
old_bits_v
VName
arr'
(TPrimExp Int64 VName -> TExp Int32
forall t v. IntExp t => TPrimExp t v -> TPrimExp Int32 v
sExt32 (TPrimExp Int64 VName -> TExp Int32)
-> Count Elements (TPrimExp Int64 VName)
-> Count Elements (TExp Int32)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Count Elements (TPrimExp Int64 VName)
bucket_offset)
(TV Int32 -> VName
forall t. TV t -> VName
tvVar TV Int32
run_loop)
(Exp -> Exp
forall {v}. PrimExp v -> PrimExp v
toBits (VName -> PrimType -> Exp
Imp.var VName
x PrimType
t))
VName
old VName -> Exp -> ImpM MCMem HostEnv Multicore ()
forall rep r op. VName -> Exp -> ImpM rep r op ()
<~~ Exp -> Exp
forall {v}. PrimExp v -> PrimExp v
fromBits Exp
old_bits
supportedPrims :: Int -> Bool
supportedPrims :: Int -> Bool
supportedPrims Int
8 = Bool
True
supportedPrims Int
16 = Bool
True
supportedPrims Int
32 = Bool
True
supportedPrims Int
64 = Bool
True
supportedPrims Int
_ = Bool
False
toIntegral :: Int -> MulticoreGen PrimType
toIntegral :: Int -> MulticoreGen PrimType
toIntegral Int
8 = PrimType -> MulticoreGen PrimType
forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
int8
toIntegral Int
16 = PrimType -> MulticoreGen PrimType
forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
int16
toIntegral Int
32 = PrimType -> MulticoreGen PrimType
forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
int32
toIntegral Int
64 = PrimType -> MulticoreGen PrimType
forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
int64
toIntegral Int
b = String -> MulticoreGen PrimType
forall a. HasCallStack => String -> a
error (String -> MulticoreGen PrimType)
-> String -> MulticoreGen PrimType
forall a b. (a -> b) -> a -> b
$ String
"number of bytes is not supported for CAS - " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Pretty a => a -> String
pretty Int
b