module Data.Array.Parallel.Unlifted.Distributed.Primitive.DistST
( DistST
, stToDistST
, distST_, distST
, runDistST, runDistST_seq
, myIndex
, myD
, readMyMD, writeMyMD
, mapDST_, mapDST, zipWithDST_, zipWithDST)
where
import qualified Data.Array.Parallel.Unlifted.Distributed.What as W
import Data.Array.Parallel.Unlifted.Distributed.Primitive.DT
import Data.Array.Parallel.Unlifted.Distributed.Primitive.Gang
import Data.Array.Parallel.Unlifted.Distributed.Data.Tuple
import Data.Array.Parallel.Base (ST, runST)
import Control.Monad (liftM)
newtype DistST s a = DistST { unDistST :: Int -> ST s a }
instance Monad (DistST s) where
return = DistST . const . return
DistST p >>= f = DistST $ \i -> do
x <- p i
unDistST (f x) i
myIndex :: DistST s Int
myIndex = DistST return
stToDistST :: ST s a -> DistST s a
stToDistST p = DistST $ \_ -> p
myD :: DT a => Dist a -> DistST s a
myD dt = liftM (indexD "myD" dt) myIndex
readMyMD :: DT a => MDist a s -> DistST s a
readMyMD mdt
= do i <- myIndex
stToDistST $ readMD mdt i
writeMyMD :: DT a => MDist a s -> a -> DistST s ()
writeMyMD mdt x
= do i <- myIndex
stToDistST $ writeMD mdt i x
runDistST :: DT a => W.Comp -> Gang -> (forall s. DistST s a) -> Dist a
runDistST comp g p
= runST $ distST comp g p
runDistST_seq
:: forall a. DT a
=> Gang -> (forall s. DistST s a) -> Dist a
runDistST_seq g p
= runST
$ do
md <- newMD g
go md 0
unsafeFreezeMD md
where
!n = gangSize g
go :: forall s. MDist a s -> Int -> ST s ()
go md i | i < n = do
writeMD md i =<< unDistST p i
go md (i+1)
| otherwise = return ()
distST :: DT a
=> W.Comp -> Gang
-> DistST s a -> ST s (Dist a)
distST comp g p
= do md <- newMD g
distST_ comp g
$ writeMyMD md =<< p
unsafeFreezeMD md
distST_ :: W.Comp -> Gang -> DistST s () -> ST s ()
distST_ comp gang proc
= gangST gang
(show comp)
(workloadOfComp comp)
$ unDistST proc
workloadOfComp :: W.Comp -> Workload
workloadOfComp cc
= case cc of
W.CDist w -> workloadOfWhat w
_ -> WorkUnknown
workloadOfWhat :: W.What -> Workload
workloadOfWhat ww
= case ww of
W.WJoinCopy elems -> WorkCopy elems
_ -> WorkUnknown
mapDST :: (DT a, DT b)
=> W.What -> Gang -> (a -> DistST s b) -> Dist a -> ST s (Dist b)
mapDST what g p !d
= mapDST' what g (\x -> x `deepSeqD` p x) d
mapDST_ :: DT a => W.What -> Gang -> (a -> DistST s ()) -> Dist a -> ST s ()
mapDST_ what g p !d
= mapDST_' what g (\x -> x `deepSeqD` p x) d
mapDST' :: (DT a, DT b) => W.What -> Gang -> (a -> DistST s b) -> Dist a -> ST s (Dist b)
mapDST' what g p !d
= distST (W.CDist what) g (myD d >>= p)
mapDST_'
:: DT a
=> W.What -> Gang -> (a -> DistST s ()) -> Dist a -> ST s ()
mapDST_' what g p !d
= distST_ (W.CDist what) g (myD d >>= p)
zipWithDST
:: (DT a, DT b, DT c)
=> W.What
-> Gang
-> (a -> b -> DistST s c) -> Dist a -> Dist b -> ST s (Dist c)
zipWithDST what g p !dx !dy
= mapDST what g (uncurry p) (zipD dx dy)
zipWithDST_
:: (DT a, DT b)
=> W.What -> Gang -> (a -> b -> DistST s ()) -> Dist a -> Dist b -> ST s ()
zipWithDST_ what g p !dx !dy
= mapDST_ what g (uncurry p) (zipD dx dy)