Add combine2ByTagUP
[packages/dph.git] / dph-prim-par / Data / Array / Parallel / Unlifted / Parallel / Combinators.hs
1 -----------------------------------------------------------------------------
2 -- |
3 -- Module : Data.Array.Parallel.Unlifted.Parallel.Combinators
4 -- Copyright : (c) 2006 Roman Leshchinskiy
5 -- License : see libraries/ndp/LICENSE
6 --
7 -- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au>
8 -- Stability : experimental
9 -- Portability : portable
10 --
11 -- Description ---------------------------------------------------------------
12 --
13 -- Parallel combinators for unlifted arrays
14 --
15
16 {-# LANGUAGE CPP #-}
17
18 #include "fusion-phases.h"
19
20 module Data.Array.Parallel.Unlifted.Parallel.Combinators (
21 mapUP, filterUP, packUP, combineUP, combine2ByTagUP,
22 zipWithUP, foldUP, fold1UP, foldl1UP, scanUP
23 ) where
24
25 import Data.Array.Parallel.Base
26 import Data.Array.Parallel.Unlifted.Sequential
27 import Data.Array.Parallel.Unlifted.Distributed
28
29 mapUP :: (UA a, UA b) => (a -> b) -> UArr a -> UArr b
30 {-# INLINE mapUP #-}
31 mapUP f = splitJoinD theGang (mapD theGang (mapU f))
32
33 filterUP :: UA a => (a -> Bool) -> UArr a -> UArr a
34 {-# INLINE filterUP #-}
35 filterUP f = joinD theGang unbalanced
36 . mapD theGang (filterU f)
37 . splitD theGang unbalanced
38
39
40
41 -- |Extract all elements from an array according to a given flag array
42 --
43 packUP:: UA e => UArr e -> UArr Bool -> UArr e
44 {-# INLINE_UP packUP #-}
45 packUP xs flags = fstU . filterUP sndS $ zipU xs flags
46
47 combineUP :: UA a => UArr Bool -> UArr a -> UArr a -> UArr a
48 {-# INLINE_UP combineUP #-}
49 combineUP flags !xs !ys = joinD theGang balanced
50 . zipWithD theGang go (zipD is ns)
51 $ splitD theGang balanced flags
52 where
53 ns = mapD theGang count
54 $ splitD theGang balanced flags
55
56 is = fstS $ scanD theGang add (0 :*: 0) ns
57
58 count bs = let ts = sumU (mapU fromBool bs)
59 in ts :*: (lengthU bs - ts)
60
61 add (x1 :*: y1) (x2 :*: y2) = (x1 + x2) :*: (y1 + y2)
62
63 go ((i :*: j) :*: (m :*: n)) bs = combineU bs (sliceU xs i m) (sliceU ys j n)
64
65 combine2ByTagUP :: UA a => UArr Int -> UArr a -> UArr a -> UArr a
66 {-# INLINE_UP combine2ByTagUP #-}
67 combine2ByTagUP tags !xs !ys = joinD theGang balanced
68 $ zipWithD theGang go (zipD is ns)
69 $ splitD theGang balanced tags
70 where
71 ns = mapD theGang count
72 $ splitD theGang balanced tags
73
74 count bs = let ones = sumU bs
75 in (lengthU bs - ones) :*: ones
76
77 is = fstS $ scanD theGang add (0 :*: 0) ns
78
79 add (x1 :*: y1) (x2 :*: y2) = (x1+x2) :*: (y1+y2)
80
81 go ((i :*: j) :*: (m :*: n)) ts = combine2ByTagU ts (sliceU xs i m)
82 (sliceU ys j n)
83
84
85 zipWithUP :: (UA a, UA b, UA c) => (a -> b -> c) -> UArr a -> UArr b -> UArr c
86 {-# INLINE zipWithUP #-}
87 zipWithUP f xs ys = splitJoinD theGang (mapD theGang (mapU (uncurryS f))) (zipU xs ys)
88 {-
89 zipWithUP f a b = joinD theGang balanced
90 (zipWithD theGang (zipWithU f)
91 (splitD theGang balanced a)
92 (splitD theGang balanced b))
93 -}
94 --zipWithUP f a b = mapUP (uncurryS f) (zipU a b)
95
96 foldUP :: (UA a, DT a) => (a -> a -> a) -> a -> UArr a -> a
97 {-# INLINE foldUP #-}
98 foldUP f z xs = maybeS z (f z)
99 (foldD theGang combine
100 (mapD theGang (foldl1MaybeU f)
101 (splitD theGang unbalanced
102 xs)))
103 where
104 combine (JustS x) (JustS y) = JustS (f x y)
105 combine (JustS x) NothingS = JustS x
106 combine NothingS (JustS y) = JustS y
107 combine NothingS NothingS = NothingS
108
109
110 -- |Array reduction proceeding from the left (requires associative combination)
111 --
112 foldlUP :: (DT a, UA a) => (a -> a -> a) -> a -> UArr a -> a
113 {-# INLINE_UP foldlUP #-}
114 foldlUP f z arr
115 | nullU arr = z
116 | otherwise = foldl1UP f arr
117
118 -- |Reduction of a non-empty array which requires an associative combination
119 -- function
120 --
121 fold1UP :: (DT a, UA a) => (a -> a -> a) -> UArr a -> a
122 {-# INLINE fold1UP #-}
123 fold1UP = foldl1UP
124
125
126
127 foldl1UP :: (DT a, UA a) => (a -> a -> a) -> UArr a -> a
128 {-# INLINE_U foldl1UP #-}
129 foldl1UP f arr = (maybeS z (f z)
130 . foldD theGang combine
131 . mapD theGang (foldl1MaybeU f)
132 . splitD theGang unbalanced) arr
133 where
134 z = arr !: 0
135 combine (JustS x) (JustS y) = JustS (f x y)
136 combine (JustS x) NothingS = JustS x
137 combine NothingS (JustS y) = JustS y
138 combine NothingS NothingS = NothingS
139
140 scanUP :: (DT a, UA a) => (a -> a -> a) -> a -> UArr a -> UArr a
141 {-# INLINE_UP scanUP #-}
142 scanUP f z = splitJoinD theGang go
143 where
144 go xs = let ds :*: zs = unzipD $ mapD theGang (scanResU f z) xs
145 zs' = fstS (scanD theGang f z zs)
146 in
147 zipWithD theGang (mapU . f) zs' ds
148
149