{-# LANGUAGE BangPatterns #-}
module Control.Monad.Par.Combinator
(
parMap, parMapM,
parMapReduceRangeThresh, parMapReduceRange,
InclusiveRange(..),
parFor
)
where
import Control.DeepSeq
import Data.Traversable
import Control.Monad as M hiding (mapM, sequence, join)
import Prelude hiding (mapM, sequence, head,tail)
import GHC.Conc (numCapabilities)
import Control.Monad.Par.Class
parMap :: (Traversable t, NFData b, ParFuture iv p) => (a -> b) -> t a -> p (t b)
parMap :: forall (t :: * -> *) b (iv :: * -> *) (p :: * -> *) a.
(Traversable t, NFData b, ParFuture iv p) =>
(a -> b) -> t a -> p (t b)
parMap a -> b
f t a
xs = (a -> p (iv b)) -> t a -> p (t (iv b))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> t a -> m (t b)
mapM (b -> p (iv b)
forall a. NFData a => a -> p (iv a)
forall (future :: * -> *) (m :: * -> *) a.
(ParFuture future m, NFData a) =>
a -> m (future a)
spawnP (b -> p (iv b)) -> (a -> b) -> a -> p (iv b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f) t a
xs p (t (iv b)) -> (t (iv b) -> p (t b)) -> p (t b)
forall a b. p a -> (a -> p b) -> p b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (iv b -> p b) -> t (iv b) -> p (t b)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> t a -> m (t b)
mapM iv b -> p b
forall a. iv a -> p a
forall (future :: * -> *) (m :: * -> *) a.
ParFuture future m =>
future a -> m a
get
parMapM :: (Traversable t, NFData b, ParFuture iv p) => (a -> p b) -> t a -> p (t b)
parMapM :: forall (t :: * -> *) b (iv :: * -> *) (p :: * -> *) a.
(Traversable t, NFData b, ParFuture iv p) =>
(a -> p b) -> t a -> p (t b)
parMapM a -> p b
f t a
xs = (a -> p (iv b)) -> t a -> p (t (iv b))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> t a -> m (t b)
mapM (p b -> p (iv b)
forall a. NFData a => p a -> p (iv a)
forall (future :: * -> *) (m :: * -> *) a.
(ParFuture future m, NFData a) =>
m a -> m (future a)
spawn (p b -> p (iv b)) -> (a -> p b) -> a -> p (iv b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> p b
f) t a
xs p (t (iv b)) -> (t (iv b) -> p (t b)) -> p (t b)
forall a b. p a -> (a -> p b) -> p b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (iv b -> p b) -> t (iv b) -> p (t b)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> t a -> m (t b)
mapM iv b -> p b
forall a. iv a -> p a
forall (future :: * -> *) (m :: * -> *) a.
ParFuture future m =>
future a -> m a
get
data InclusiveRange = InclusiveRange Int Int
parMapReduceRangeThresh
:: (NFData a, ParFuture iv p)
=> Int
-> InclusiveRange
-> (Int -> p a)
-> (a -> a -> p a)
-> a
-> p a
parMapReduceRangeThresh :: forall a (iv :: * -> *) (p :: * -> *).
(NFData a, ParFuture iv p) =>
Int
-> InclusiveRange -> (Int -> p a) -> (a -> a -> p a) -> a -> p a
parMapReduceRangeThresh Int
threshold (InclusiveRange Int
min Int
max) Int -> p a
fn a -> a -> p a
binop a
init
= Int -> Int -> p a
loop Int
min Int
max
where
loop :: Int -> Int -> p a
loop Int
min Int
max
| Int
max Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
min Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
threshold =
let mapred :: a -> Int -> p a
mapred a
a Int
b = do a
x <- Int -> p a
fn Int
b;
a
result <- a
a a -> a -> p a
`binop` a
x
a -> p a
forall a. a -> p a
forall (m :: * -> *) a. Monad m => a -> m a
return a
result
in (a -> Int -> p a) -> a -> [Int] -> p a
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM a -> Int -> p a
mapred a
init [Int
min..Int
max]
| Bool
otherwise = do
let mid :: Int
mid = Int
min Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ((Int
max Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
min) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
2)
iv a
rght <- p a -> p (iv a)
forall a. NFData a => p a -> p (iv a)
forall (future :: * -> *) (m :: * -> *) a.
(ParFuture future m, NFData a) =>
m a -> m (future a)
spawn (p a -> p (iv a)) -> p a -> p (iv a)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> p a
loop (Int
midInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
max
a
l <- Int -> Int -> p a
loop Int
min Int
mid
a
r <- iv a -> p a
forall a. iv a -> p a
forall (future :: * -> *) (m :: * -> *) a.
ParFuture future m =>
future a -> m a
get iv a
rght
a
l a -> a -> p a
`binop` a
r
auto_partition_factor :: Int
auto_partition_factor :: Int
auto_partition_factor = Int
4
parMapReduceRange :: (NFData a, ParFuture iv p) =>
InclusiveRange -> (Int -> p a) -> (a -> a -> p a) -> a -> p a
parMapReduceRange :: forall a (iv :: * -> *) (p :: * -> *).
(NFData a, ParFuture iv p) =>
InclusiveRange -> (Int -> p a) -> (a -> a -> p a) -> a -> p a
parMapReduceRange (InclusiveRange Int
start Int
end) Int -> p a
fn a -> a -> p a
binop a
init =
Int -> [(Int, Int)] -> p a
loop ([(Int, Int)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, Int)]
segs) [(Int, Int)]
segs
where
segs :: [(Int, Int)]
segs = Int -> (Int, Int) -> [(Int, Int)]
splitInclusiveRange (Int
auto_partition_factor Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
numCapabilities) (Int
start,Int
end)
loop :: Int -> [(Int, Int)] -> p a
loop Int
1 [(Int
st,Int
en)] =
let mapred :: a -> Int -> p a
mapred a
a Int
b = do a
x <- Int -> p a
fn Int
b;
a
result <- a
a a -> a -> p a
`binop` a
x
a -> p a
forall a. a -> p a
forall (m :: * -> *) a. Monad m => a -> m a
return a
result
in (a -> Int -> p a) -> a -> [Int] -> p a
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM a -> Int -> p a
mapred a
init [Int
st..Int
en]
loop Int
n [(Int, Int)]
segs =
let half :: Int
half = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
2
([(Int, Int)]
left,[(Int, Int)]
right) = Int -> [(Int, Int)] -> ([(Int, Int)], [(Int, Int)])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
half [(Int, Int)]
segs in
do iv a
l <- p a -> p (iv a)
forall a. NFData a => p a -> p (iv a)
forall (future :: * -> *) (m :: * -> *) a.
(ParFuture future m, NFData a) =>
m a -> m (future a)
spawn(p a -> p (iv a)) -> p a -> p (iv a)
forall a b. (a -> b) -> a -> b
$ Int -> [(Int, Int)] -> p a
loop Int
half [(Int, Int)]
left
a
r <- Int -> [(Int, Int)] -> p a
loop (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
half) [(Int, Int)]
right
a
l' <- iv a -> p a
forall a. iv a -> p a
forall (future :: * -> *) (m :: * -> *) a.
ParFuture future m =>
future a -> m a
get iv a
l
a
l' a -> a -> p a
`binop` a
r
parFor :: (ParFuture iv p) => InclusiveRange -> (Int -> p ()) -> p ()
parFor :: forall (iv :: * -> *) (p :: * -> *).
ParFuture iv p =>
InclusiveRange -> (Int -> p ()) -> p ()
parFor (InclusiveRange Int
start Int
end) Int -> p ()
body =
do
let run :: (Int, Int) -> p ()
run (Int
x,Int
y) = Int -> Int -> (Int -> p ()) -> p ()
forall (m :: * -> *).
Monad m =>
Int -> Int -> (Int -> m ()) -> m ()
for_ Int
x (Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int -> p ()
body
range_segments :: [(Int, Int)]
range_segments = Int -> (Int, Int) -> [(Int, Int)]
splitInclusiveRange (Int
4Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
numCapabilities) (Int
start,Int
end)
[iv ()]
vars <- [(Int, Int)] -> ((Int, Int) -> p (iv ())) -> p [iv ()]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
M.forM [(Int, Int)]
range_segments (\ (Int, Int)
pr -> p () -> p (iv ())
forall a. p a -> p (iv a)
forall (future :: * -> *) (m :: * -> *) a.
ParFuture future m =>
m a -> m (future a)
spawn_ ((Int, Int) -> p ()
run (Int, Int)
pr))
(iv () -> p ()) -> [iv ()] -> p ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
M.mapM_ iv () -> p ()
forall a. iv a -> p a
forall (future :: * -> *) (m :: * -> *) a.
ParFuture future m =>
future a -> m a
get [iv ()]
vars
() -> p ()
forall a. a -> p a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
splitInclusiveRange :: Int -> (Int, Int) -> [(Int, Int)]
splitInclusiveRange :: Int -> (Int, Int) -> [(Int, Int)]
splitInclusiveRange Int
pieces (Int
start,Int
end) =
(Int -> (Int, Int)) -> [Int] -> [(Int, Int)]
forall a b. (a -> b) -> [a] -> [b]
map Int -> (Int, Int)
largepiece [Int
0..Int
remainInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] [(Int, Int)] -> [(Int, Int)] -> [(Int, Int)]
forall a. [a] -> [a] -> [a]
++
(Int -> (Int, Int)) -> [Int] -> [(Int, Int)]
forall a b. (a -> b) -> [a] -> [b]
map Int -> (Int, Int)
smallpiece [Int
remain..Int
piecesInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]
where
len :: Int
len = Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
(Int
portion, Int
remain) = Int
len Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
pieces
largepiece :: Int -> (Int, Int)
largepiece Int
i =
let offset :: Int
offset = Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
portion Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
in (Int
offset, Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
portion)
smallpiece :: Int -> (Int, Int)
smallpiece Int
i =
let offset :: Int
offset = Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
portion) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
remain
in (Int
offset, Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
portion Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
{-# INLINE for_ #-}
for_ :: Monad m => Int -> Int -> (Int -> m ()) -> m ()
for_ :: forall (m :: * -> *).
Monad m =>
Int -> Int -> (Int -> m ()) -> m ()
for_ Int
start Int
end Int -> m ()
_fn | Int
start Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
end = [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"for_: start is greater than end"
for_ Int
start Int
end Int -> m ()
fn = Int -> m ()
loop Int
start
where
loop :: Int -> m ()
loop !Int
i | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
end = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
| Bool
otherwise = do Int -> m ()
fn Int
i; Int -> m ()
loop (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)