Prepare dph for a vectInfoVar type change
[packages/dph.git] / dph-prim-par / Data / Array / Parallel / Unlifted / Parallel / Combinators.hs
1 {-# LANGUAGE CPP #-}
2 #include "fusion-phases.h"
3
4 -- | Parallel combinators for unlifted arrays.
5 module Data.Array.Parallel.Unlifted.Parallel.Combinators
6 ( mapUP
7 , filterUP
8 , packUP
9 , combineUP, combine2UP
10 , zipWithUP
11 , foldUP, foldlUP, fold1UP, foldl1UP
12 , scanUP)
13 where
14 import Data.Array.Parallel.Base
15 import Data.Array.Parallel.Unlifted.Distributed
16 import Data.Array.Parallel.Unlifted.Distributed.What
17 import Data.Array.Parallel.Unlifted.Parallel.UPSel
18 import Data.Array.Parallel.Unlifted.Sequential.Vector as Seq
19
20 here :: String -> String
21 here s = "Data.Array.Parallel.Unlifted.Parallel.Combinators." Prelude.++ s
22
23
24 -- | Apply a worker to all elements of an array.
25 mapUP :: (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
26 mapUP f xs
27 = splitJoinD theGang
28 (mapD (What "mapUP/map") theGang (Seq.map f)) xs
29 {-# INLINE_UP mapUP #-}
30
31
32 -- | Keep elements that match the given predicate.
33 filterUP :: Unbox a => (a -> Bool) -> Vector a -> Vector a
34 filterUP f
35 = joinD theGang unbalanced
36 . mapD (What "filterUP/filter") theGang (Seq.filter f)
37 . splitD theGang unbalanced
38 {-# INLINE_UP filterUP #-}
39
40
41 -- | Take elements of an array where a flag value is true, and pack them into
42 -- the result.
43 --
44 -- * The souce and flag arrays must have the same length, but this is not checked.
45 --
46 packUP :: Unbox e => Vector e -> Vector Bool -> Vector e
47 packUP xs flags
48 = Seq.fsts . filterUP snd $ Seq.zip xs flags
49 {-# INLINE_UP packUP #-}
50
51
52 -- | Combine two vectors based on a selector.
53 -- If the selector is true then take the element from the first vector,
54 -- otherwise take it from the second.
55 --
56 -- * The data vectors must have enough elements to satisfy the flag vector,
57 -- but this is not checked.
58 --
59 combineUP :: Unbox a => Vector Bool -> Vector a -> Vector a -> Vector a
60 combineUP flags xs ys
61 = checkEq (here "combineUP")
62 ("tags length /= sum of args length")
63 (Seq.length flags) (Seq.length xs + Seq.length ys)
64 $ combine2UP tags (mkUPSelRep2 tags) xs ys
65 where tags = Seq.map (fromBool . not) flags
66 {-# INLINE combineUP #-}
67
68
69 -- | Combine two vectors based on a selector.
70 --
71 -- * The data vectors must have enough elements to satisfy the selector,
72 -- but this is not checked.
73 --
74 combine2UP :: Unbox a => Vector Tag -> UPSelRep2 -> Vector a -> Vector a -> Vector a
75 combine2UP tags rep !xs !ys
76 = checkEq (here "combine2UP")
77 ("tags length /= sum of args length")
78 (Seq.length tags) (Seq.length xs + Seq.length ys)
79 $ joinD theGang balanced
80 $ zipWithD (What "combine2UP/go") theGang go rep
81 $ splitD theGang balanced tags
82 where go ((i,j), (m,n)) ts
83 = Seq.combine2ByTag ts
84 (Seq.slice (here "combine2UP") xs i m)
85 (Seq.slice (here "combine2UP") ys j n)
86 {-# INLINE_UP combine2UP #-}
87
88
89 -- | Apply a worker function to correponding elements of two arrays.
90 zipWithUP :: (Unbox a, Unbox b, Unbox c)
91 => (a -> b -> c) -> Vector a -> Vector b -> Vector c
92 zipWithUP f xs ys
93 = splitJoinD theGang
94 (mapD (What "zipWithUP/map") theGang (Seq.map (uncurry f)))
95 (Seq.zip xs ys)
96 {-# INLINE_UP zipWithUP #-}
97
98
99 -- | Undirected fold.
100 -- Note that this function has more constraints on its parameters than the
101 -- standard fold function from the Haskell Prelude.
102 --
103 -- * The worker function must be associative.
104 --
105 -- * The provided starting element must be neutral with respect to the worker.
106 -- For example 0 is neutral wrt (+) and 1 is neutral wrt (*).
107 --
108 -- We need these constraints so that we can partition the fold across
109 -- several threads. Each thread folds a chunk of the input vector,
110 -- then we fold together all the results in the main thread.
111 --
112 foldUP :: (Unbox a, DT a) => (a -> a -> a) -> a -> Vector a -> a
113 foldUP f !z xs
114 = foldD (What "foldUP/f") theGang f
115 $ mapD (What "foldUP/fold") theGang (Seq.fold f z)
116 $ splitD theGang unbalanced xs
117 {-# INLINE_UP foldUP #-}
118
119
120 -- | Left fold over an array.
121 --
122 -- * If the vector is empty then this returns the provided neural element.
123 --
124 -- * The worker function must be associative.
125 --
126 -- * The provided starting element must be neutral with respect to the worker,
127 -- see `foldUP` for discussion.
128 --
129 foldlUP :: (DT a, Unbox a) => (a -> a -> a) -> a -> Vector a -> a
130 foldlUP f z arr
131 | Seq.null arr = z
132 | otherwise = foldl1UP f arr
133 {-# INLINE_UP foldlUP #-}
134
135
136 -- | Alias for `foldl1UP`
137 fold1UP :: (DT a, Unbox a) => (a -> a -> a) -> Vector a -> a
138 fold1UP = foldl1UP
139 {-# INLINE_UP fold1UP #-}
140
141
142 -- | Left fold over an array, using the first element of the vector as the
143 -- neural element.
144 --
145 -- * If the vector contains no elements then you'll get a bounds-check error.
146 --
147 -- * The worker function must be associative.
148 --
149 -- * The provided starting element must be neutral with respect to the worker,
150 -- see `foldUP` for discussion.
151 --
152 foldl1UP :: (DT a, Unbox a) => (a -> a -> a) -> Vector a -> a
153 foldl1UP f arr
154 = (maybe z (f z)
155 . foldD (What "fold1UP/foldD") theGang combine'
156 . mapD (What "fold1UP/fold1Maybe") theGang (Seq.foldl1Maybe f)
157 . splitD theGang unbalanced) arr
158 where
159 z = Seq.index (here "fold1UP") arr 0
160 combine' (Just x) (Just y) = Just (f x y)
161 combine' (Just x) Nothing = Just x
162 combine' Nothing (Just y) = Just y
163 combine' Nothing Nothing = Nothing
164 {-# INLINE_UP foldl1UP #-}
165
166
167 -- | Prefix scan. Similar to fold, but produce an array of the intermediate states.
168 --
169 -- * The worker function must be associative.
170 --
171 -- * The provided starting element must be neutral with respect to the worker,
172 -- see `foldUP` for discussion.
173 --
174 scanUP :: (DT a, Unbox a) => (a -> a -> a) -> a -> Vector a -> Vector a
175 scanUP f z
176 = splitJoinD theGang go
177 where go xs
178 = let (ds,zs) = unzipD
179 $ mapD (What "scanUP/scanRes") theGang (Seq.scanRes f z) xs
180
181 zs' = fst
182 $ scanD (What "scanUP/scan") theGang f z zs
183
184 in zipWithD (What "scanUP/map") theGang (Seq.map . f) zs' ds
185 {-# INLINE_UP scanUP #-}
186