Make combine2 work with selectors
[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, combine2UP,
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 xs = splitJoinD theGang (mapD theGang (mapU f)) xs
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 combine2UP :: UA a => USel2 -> UArr a -> UArr a -> UArr a
66 {-# INLINE_UP combine2UP #-}
67 combine2UP sel !xs !ys = zipWithUP get (tagsUSel2 sel) (indicesUSel2 sel)
68 where
69 get 0 i = xs !: i
70 get _ i = ys !: i
71
72 {-
73 combine2UP tags !xs !ys = joinD theGang balanced
74 $ zipWithD theGang go (zipD is ns)
75 $ splitD theGang balanced tags
76 where
77 ns = mapD theGang count
78 $ splitD theGang balanced tags
79
80 count bs = let ones = sumU bs
81 in (lengthU bs - ones) :*: ones
82
83 is = fstS $ scanD theGang add (0 :*: 0) ns
84
85 add (x1 :*: y1) (x2 :*: y2) = (x1+x2) :*: (y1+y2)
86
87 go ((i :*: j) :*: (m :*: n)) ts = combine2ByTagU ts (sliceU xs i m)
88 (sliceU ys j n)
89 -}
90
91 zipWithUP :: (UA a, UA b, UA c) => (a -> b -> c) -> UArr a -> UArr b -> UArr c
92 {-# INLINE zipWithUP #-}
93 zipWithUP f xs ys = splitJoinD theGang (mapD theGang (mapU (uncurryS f))) (zipU xs ys)
94 {-
95 zipWithUP f a b = joinD theGang balanced
96 (zipWithD theGang (zipWithU f)
97 (splitD theGang balanced a)
98 (splitD theGang balanced b))
99 -}
100 --zipWithUP f a b = mapUP (uncurryS f) (zipU a b)
101
102 foldUP :: (UA a, DT a) => (a -> a -> a) -> a -> UArr a -> a
103 {-# INLINE foldUP #-}
104 foldUP f !z xs = foldD theGang f
105 (mapD theGang (foldU f z)
106 (splitD theGang unbalanced xs))
107 {-
108 foldUP f z xs = maybeS z (f z)
109 (foldD theGang combine
110 (mapD theGang (foldl1MaybeU f)
111 (splitD theGang unbalanced
112 xs)))
113 where
114 combine (JustS x) (JustS y) = JustS (f x y)
115 combine (JustS x) NothingS = JustS x
116 combine NothingS (JustS y) = JustS y
117 combine NothingS NothingS = NothingS
118 -}
119
120 -- |Array reduction proceeding from the left (requires associative combination)
121 --
122 foldlUP :: (DT a, UA a) => (a -> a -> a) -> a -> UArr a -> a
123 {-# INLINE_UP foldlUP #-}
124 foldlUP f z arr
125 | nullU arr = z
126 | otherwise = foldl1UP f arr
127
128 -- |Reduction of a non-empty array which requires an associative combination
129 -- function
130 --
131 fold1UP :: (DT a, UA a) => (a -> a -> a) -> UArr a -> a
132 {-# INLINE fold1UP #-}
133 fold1UP = foldl1UP
134
135
136
137 foldl1UP :: (DT a, UA a) => (a -> a -> a) -> UArr a -> a
138 {-# INLINE_U foldl1UP #-}
139 foldl1UP f arr = (maybeS z (f z)
140 . foldD theGang combine
141 . mapD theGang (foldl1MaybeU f)
142 . splitD theGang unbalanced) arr
143 where
144 z = arr !: 0
145 combine (JustS x) (JustS y) = JustS (f x y)
146 combine (JustS x) NothingS = JustS x
147 combine NothingS (JustS y) = JustS y
148 combine NothingS NothingS = NothingS
149
150 scanUP :: (DT a, UA a) => (a -> a -> a) -> a -> UArr a -> UArr a
151 {-# INLINE_UP scanUP #-}
152 scanUP f z = splitJoinD theGang go
153 where
154 go xs = let ds :*: zs = unzipD $ mapD theGang (scanResU f z) xs
155 zs' = fstS (scanD theGang f z zs)
156 in
157 zipWithD theGang (mapU . f) zs' ds
158
159