dph-lifted-vseg: push indexvs into lifted library
[packages/dph.git] / dph-lifted-vseg / Data / Array / Parallel / PArray / PData / Tuple2.hs
1 {-# OPTIONS_HADDOCK hide #-}
2 #include "fusion-phases.h"
3
4 -- | PR instance for tuples.
5 module Data.Array.Parallel.PArray.PData.Tuple2
6 ( PData(..), PDatas(..)
7 , zipPD
8 , ziplPR
9 , unzipPD
10 , unziplPD)
11 where
12 import Data.Array.Parallel.Pretty
13 import Data.Array.Parallel.PArray.PData.Base
14 import Data.Array.Parallel.PArray.PData.Nested
15 import GHC.Exts
16 import Prelude hiding (zip, unzip)
17 import qualified Data.Vector as V
18 import qualified Prelude as P
19 import qualified Data.Array.Parallel.Unlifted as U
20
21 -------------------------------------------------------------------------------
22 data instance PData (a, b)
23 = PTuple2 !(PData a) !(PData b)
24
25 data instance PDatas (a, b)
26 = PTuple2s !(PDatas a) !(PDatas b)
27
28
29 -- PR -------------------------------------------------------------------------
30 instance (PR a, PR b) => PR (a, b) where
31
32 {-# NOINLINE validPR #-}
33 validPR (PTuple2 xs ys)
34 = validPR xs && validPR ys
35
36
37 {-# NOINLINE nfPR #-}
38 nfPR (PTuple2 arr1 arr2)
39 = nfPR arr1 `seq` nfPR arr2 `seq` ()
40
41
42 {-# NOINLINE similarPR #-}
43 similarPR (x1, y1) (x2, y2)
44 = similarPR x1 x2
45 && similarPR y1 y2
46
47
48 {-# NOINLINE coversPR #-}
49 coversPR weak (PTuple2 arr1 arr2) ix
50 = coversPR weak arr1 ix
51 && coversPR weak arr2 ix
52
53 {-# NOINLINE pprpPR #-}
54 pprpPR (x, y)
55 = text "Tuple2 " <> vcat [pprpPR x, pprpPR y]
56
57
58 {-# NOINLINE pprpDataPR #-}
59 pprpDataPR (PTuple2 xs ys)
60 = text "PTuple2 " <> vcat [pprpDataPR xs, pprpDataPR ys]
61
62 -- Constructors -------------------------------
63 {-# INLINE_PDATA emptyPR #-}
64 emptyPR
65 = PTuple2 emptyPR emptyPR
66
67
68 {-# INLINE_PDATA replicatePR #-}
69 replicatePR len (x, y)
70 = PTuple2 (replicatePR len x)
71 (replicatePR len y)
72
73
74 {-# INLINE_PDATA replicatesPR #-}
75 replicatesPR lens (PTuple2 arr1 arr2)
76 = PTuple2 (replicatesPR lens arr1)
77 (replicatesPR lens arr2)
78
79
80 {-# INLINE_PDATA appendPR #-}
81 appendPR (PTuple2 arr11 arr12) (PTuple2 arr21 arr22)
82 = PTuple2 (arr11 `appendPR` arr21)
83 (arr12 `appendPR` arr22)
84
85
86 {-# INLINE_PDATA appendsPR #-}
87 appendsPR segdResult segd1 (PTuple2 arrs11 arrs12) segd2 (PTuple2 arrs21 arrs22)
88 = PTuple2 (appendsPR segdResult segd1 arrs11 segd2 arrs21)
89 (appendsPR segdResult segd1 arrs12 segd2 arrs22)
90
91
92 -- Projections ---------------------------------
93 {-# INLINE_PDATA lengthPR #-}
94 lengthPR (PTuple2 arr1 _)
95 = lengthPR arr1
96
97 {-# INLINE_PDATA indexPR #-}
98 indexPR (PTuple2 arr1 arr2) ix
99 = (indexPR arr1 ix, indexPR arr2 ix)
100
101 {-# INLINE_PDATA indexsPR #-}
102 indexsPR (PTuple2s xs ys) srcixs
103 = PTuple2 (indexsPR xs srcixs)
104 (indexsPR ys srcixs)
105
106 {-# INLINE_PDATA indexvsPR #-}
107 indexvsPR (PTuple2s xs ys) vsegd srcixs
108 = PTuple2 (indexvsPR xs vsegd srcixs)
109 (indexvsPR ys vsegd srcixs)
110
111 {-# INLINE_PDATA extractPR #-}
112 extractPR (PTuple2 arr1 arr2) start len
113 = PTuple2 (extractPR arr1 start len)
114 (extractPR arr2 start len)
115
116 {-# INLINE_PDATA extractssPR #-}
117 extractssPR (PTuple2s xs ys) ussegd
118 = PTuple2 (extractssPR xs ussegd)
119 (extractssPR ys ussegd)
120
121 {-# INLINE_PDATA extractvsPR #-}
122 extractvsPR (PTuple2s xs ys) uvsegd
123 = PTuple2 (extractvsPR xs uvsegd)
124 (extractvsPR ys uvsegd)
125
126
127 -- Pack and Combine ---------------------------
128 {-# INLINE_PDATA packByTagPR #-}
129 packByTagPR (PTuple2 arr1 arr2) tags tag
130 = PTuple2 (packByTagPR arr1 tags tag)
131 (packByTagPR arr2 tags tag)
132
133 {-# INLINE_PDATA combine2PR #-}
134 combine2PR sel (PTuple2 xs1 ys1) (PTuple2 xs2 ys2)
135 = PTuple2 (combine2PR sel xs1 xs2)
136 (combine2PR sel ys1 ys2)
137
138
139 -- Conversions --------------------------------
140 {-# NOINLINE fromVectorPR #-}
141 fromVectorPR vec
142 = let (xs, ys) = V.unzip vec
143 in PTuple2 (fromVectorPR xs)
144 (fromVectorPR ys)
145
146 {-# NOINLINE toVectorPR #-}
147 toVectorPR (PTuple2 xs ys)
148 = V.zip (toVectorPR xs)
149 (toVectorPR ys)
150
151
152 -- PData --------------------------------------
153 {-# INLINE_PDATA emptydPR #-}
154 emptydPR
155 = PTuple2s emptydPR
156 emptydPR
157
158
159 {-# INLINE_PDATA singletondPR #-}
160 singletondPR (PTuple2 x y)
161 = PTuple2s (singletondPR x)
162 (singletondPR y)
163
164
165 {-# INLINE_PDATA lengthdPR #-}
166 lengthdPR (PTuple2s xs _)
167 = lengthdPR xs
168
169
170 {-# INLINE_PDATA indexdPR #-}
171 indexdPR (PTuple2s xs ys) i
172 = PTuple2 (indexdPR xs i)
173 (indexdPR ys i)
174
175
176 {-# INLINE_PDATA appenddPR #-}
177 appenddPR (PTuple2s xs1 ys1) (PTuple2s xs2 ys2)
178 = PTuple2s (appenddPR xs1 xs2)
179 (appenddPR ys1 ys2)
180
181
182 {-# NOINLINE fromVectordPR #-}
183 fromVectordPR vec
184 = let (xss, yss) = V.unzip $ V.map (\(PTuple2 xs ys) -> (xs, ys)) vec
185 in PTuple2s (fromVectordPR xss)
186 (fromVectordPR yss)
187
188
189 {-# NOINLINE toVectordPR #-}
190 toVectordPR (PTuple2s pdatas1 pdatas2)
191 = V.zipWith PTuple2
192 (toVectordPR pdatas1)
193 (toVectordPR pdatas2)
194
195
196 -- PD Functions ---------------------------------------------------------------
197 -- These work on PData arrays of tuples, but don't need a PA or PR dictionary
198
199 -- | O(1). Zip a pair of arrays into an array of pairs.
200 zipPD :: PData a -> PData b -> PData (a, b)
201 zipPD = PTuple2
202 {-# INLINE_PA zipPD #-}
203
204
205 -- | Lifted zip.
206 ziplPR :: (PR a, PR b) => PData (PArray a) -> PData (PArray b) -> PData (PArray (a, b))
207 ziplPR arr1 arr2
208 = let -- We need to flatten the data here because we can't guarantee
209 -- that the vsegds of both arrays have the same form.
210 -- One of the arrays may have been created with replicate, and
211 -- thus has internal sharing, while the other does not.
212 (segd1, pdata1) = flattenPR arr1
213 (_, pdata2) = flattenPR arr2
214
215 in PNested (U.promoteSegdToVSegd segd1)
216 (PTuple2s (singletondPR pdata1) (singletondPR pdata2))
217
218 {-# INLINE_PA ziplPR #-}
219
220
221 -- | O(1). Unzip an array of pairs into a pair of arrays.
222 unzipPD :: PData (a, b) -> (PData a, PData b)
223 unzipPD (PTuple2 xs ys) = (xs, ys)
224 {-# INLINE_PA unzipPD #-}
225
226
227 -- | Lifted unzip.
228 unziplPD :: PData (PArray (a, b)) -> PData (PArray a, PArray b)
229 unziplPD (PNested uvsegd (PTuple2s xsdata ysdata))
230 = PTuple2 (PNested uvsegd xsdata)
231 (PNested uvsegd ysdata)
232 {-# INLINE_PA unziplPD #-}
233
234
235 -- Show -----------------------------------------------------------------------
236 deriving instance (Show (PData a), Show (PData b)) => Show (PData (a, b))
237 deriving instance (Show (PDatas a), Show (PDatas b)) => Show (PDatas (a, b))
238
239
240 instance ( PR a, PR b, Show a, Show b
241 , PprVirtual (PData a), PprVirtual (PData b))
242 => PprVirtual (PData (a, b)) where
243 pprv (PTuple2 xs ys)
244 = text $ show
245 $ P.zip (V.toList $ toVectorPR xs)
246 (V.toList $ toVectorPR ys)
247