44c42b19f380e8485da556491ded3d2bb08cb525
[packages/dph.git] / dph-prim-par / Data / Array / Parallel / Unlifted / Parallel / UPVSegd.hs
1 {-# LANGUAGE CPP #-}
2 #include "fusion-phases.h"
3
4 {-# OPTIONS -Wall -fno-warn-orphans -fno-warn-missing-signatures #-}
5
6 -- | Parallel virtual segment descriptors.
7 module Data.Array.Parallel.Unlifted.Parallel.UPVSegd (
8 -- * Types
9 UPVSegd,
10
11 -- * Consistency check
12 valid,
13
14 -- * Constructors
15 mkUPVSegd,
16 fromUPSegd,
17 fromUPSSegd,
18 empty,
19 singleton,
20
21 -- * Predicates
22 isManifest,
23 isContiguous,
24
25 -- * Projections
26 length,
27 takeVSegids,
28 takeUPSSegd,
29 takeLengths,
30 getSeg,
31
32 -- * Demotion
33 demoteToUPSSegd,
34 unsafeDemoteToUPSegd,
35
36 -- * Operators
37 updateVSegs,
38 appendWith,
39 combine2,
40 ) where
41 import Data.Array.Parallel.Unlifted.Parallel.Permute
42 import Data.Array.Parallel.Unlifted.Parallel.UPSel (UPSel2)
43 import Data.Array.Parallel.Unlifted.Parallel.UPSSegd (UPSSegd)
44 import Data.Array.Parallel.Unlifted.Parallel.UPSegd (UPSegd)
45 import Data.Array.Parallel.Unlifted.Sequential.Vector (Vector)
46 import Data.Array.Parallel.Pretty hiding (empty)
47 import Prelude hiding (length)
48
49 import qualified Data.Array.Parallel.Unlifted.Sequential.USSegd as USSegd
50 import qualified Data.Array.Parallel.Unlifted.Sequential.Vector as V
51 import qualified Data.Array.Parallel.Unlifted.Parallel.UPSel as UPSel
52 import qualified Data.Array.Parallel.Unlifted.Parallel.UPSegd as UPSegd
53 import qualified Data.Array.Parallel.Unlifted.Parallel.UPSSegd as UPSSegd
54
55
56 -- UPVSegd ---------------------------------------------------------------------
57 -- | A parallel virtual segment descriptor is an extension of `UPSSegd`
58 -- that explicitly represents sharing of data between multiple segments.
59 --
60 -- TODO: It would probably be better to represent the vsegids as a lens (function)
61 -- instead of a vector of segids. Much of the time the vsegids are just [0..n]
62 --
63 data UPVSegd
64 = UPVSegd
65 { upvsegd_manifest :: !Bool
66 -- ^ When the vsegids field holds a lazy (V.enumFromTo 0 (len - 1))
67 -- then this field is True. This lets us perform some operations like
68 -- demoteToUPSSegd without actually creating it.
69
70 , upvsegd_vsegids :: Vector Int
71 -- ^ Virtual segment identifiers that indicate what physical segment
72 -- to use for each virtual segment.
73 --
74 -- IMPORTANT:
75 --- This field must be lazy (no bang) because when it has the value
76 -- (V.enumFromTo 0 (len - 1)) we want to avoid building the enumeration
77 -- unless it's strictly demanded.
78
79 , upvsegd_upssegd :: !UPSSegd }
80 deriving (Show)
81
82
83 -- | Pretty print the physical representation of a `UVSegd`
84 instance PprPhysical UPVSegd where
85 pprp (UPVSegd _ vsegids upssegd)
86 = vcat
87 [ text "UPVSegd" $$ (nest 7 $ text "vsegids: " <+> (text $ show $ V.toList vsegids))
88 , pprp upssegd ]
89
90
91 -- | O(1). Check the internal consistency of a virutal segmentation descriptor.
92 --
93 -- * TODO: this doesn't do any checks yet.
94 --\b
95 valid :: UPVSegd -> Bool
96 valid UPVSegd{} = True
97 {-# NOINLINE valid #-}
98 -- NOINLINE because it's only used during debugging anyway.
99
100
101 -- Constructors ---------------------------------------------------------------
102 -- | O(1). Construct a new virtual segment descriptor.
103 mkUPVSegd
104 :: Vector Int -- ^ Array saying which physical segment to use for each virtual segment.
105 -> UPSSegd -- ^ Scattered segment descriptor defining the physical segments.
106 -> UPVSegd
107
108 mkUPVSegd = UPVSegd False
109 {-# NOINLINE mkUPVSegd #-}
110
111
112 -- | O(segs). Promote a `UPSSegd` to a `UPVSegd`.
113 -- The result contains one virtual segment for every physical segment
114 -- defined by the `UPSSegd`.
115 --
116 -- TODO: make this parallel, use parallel version of enumFromTo.
117 --
118 fromUPSSegd :: UPSSegd -> UPVSegd
119 fromUPSSegd upssegd
120 = UPVSegd True
121 (V.enumFromTo 0 (UPSSegd.length upssegd - 1))
122 upssegd
123 {-# NOINLINE fromUPSSegd #-}
124
125
126 -- | O(segs). Promote a `UPSegd` to a `UPVSegd`.
127 -- All segments are assumed to come from a flat array with sourceid 0.
128 -- The result contains one virtual segment for every physical segment
129 -- the provided `UPSegd`.
130 --
131 fromUPSegd :: UPSegd -> UPVSegd
132 fromUPSegd = fromUPSSegd . UPSSegd.fromUPSegd
133 {-# NOINLINE fromUPSegd #-}
134
135
136 -- | O(1). Yield an empty segment descriptor, with no elements or segments.
137 empty :: UPVSegd
138 empty = UPVSegd True V.empty UPSSegd.empty
139 {-# NOINLINE empty #-}
140
141
142 -- | O(1). Yield a singleton segment descriptor.
143 -- The single segment covers the given number of elements in a flat array
144 -- with sourceid 0.
145 singleton :: Int -> UPVSegd
146 singleton n = UPVSegd True (V.singleton 0) (UPSSegd.singleton n)
147 {-# NOINLINE singleton #-}
148
149
150 -- Predicates -----------------------------------------------------------------
151 -- | O(1). Checks whether all the segments are manifest (unshared / non-virtual).
152 -- If this is the case, then the vsegids field will be [0..len-1].
153 --
154 -- Consumers can check this field, avoid demanding the vsegids field.
155 -- This can avoid the need for it to be generated in the first place, due to
156 -- lazy evaluation.
157 --
158 isManifest :: UPVSegd -> Bool
159 isManifest = upvsegd_manifest
160 {-# INLINE isManifest #-}
161
162
163 -- | O(1). True when the starts are identical to the usegd indices field and
164 -- the sources are all 0's.
165 --
166 -- In this case all the data elements are in one contiguous flat
167 -- array, and consumers can avoid looking at the real starts and
168 -- sources fields.
169 --
170 isContiguous :: UPVSegd -> Bool
171 isContiguous = UPSSegd.isContiguous . upvsegd_upssegd
172 {-# INLINE isContiguous #-}
173
174
175 -- Projections ----------------------------------------------------------------
176 -- INLINE trivial projections as they'll expand to a single record selector.
177
178 -- | O(1). Yield the overall number of segments.
179 length :: UPVSegd -> Int
180 length = V.length . upvsegd_vsegids
181 {-# INLINE length #-}
182
183
184 -- | O(1). Yield the virtual segment ids of `UPVSegd`.
185 takeVSegids :: UPVSegd -> Vector Int
186 takeVSegids = upvsegd_vsegids
187 {-# INLINE takeVSegids #-}
188
189
190 -- | O(1). Yield the `UPSSegd` of `UPVSegd`.
191 takeUPSSegd :: UPVSegd -> UPSSegd
192 takeUPSSegd = upvsegd_upssegd
193 {-# INLINE takeUPSSegd #-}
194
195
196 -- | O(segs). Yield the lengths of the segments described by a `UPVSegd`.
197 takeLengths :: UPVSegd -> Vector Int
198 takeLengths (UPVSegd manifest vsegids upssegd)
199 | manifest = UPSSegd.takeLengths upssegd
200 | otherwise = V.map (UPSSegd.takeLengths upssegd V.!) vsegids
201 {-# NOINLINE takeLengths #-}
202
203
204 -- | O(1). Get the length, starting index, and source id of a segment.
205 --
206 -- NOTE: We don't return the segment index field from the USSegd as this refers
207 -- to the flat index relative to the SSegd array, rather than
208 -- relative to the UVSegd array. If we tried to promote the USSegd index
209 -- to a UVSegd index it could overflow.
210 --
211 getSeg :: UPVSegd -> Int -> (Int, Int, Int)
212 getSeg (UPVSegd _ vsegids upssegd) ix
213 = let (len, _index, start, source) = UPSSegd.getSeg upssegd (vsegids V.! ix)
214 in (len, start, source)
215 {-# INLINE_UP getSeg #-}
216
217
218 -- Demotion -------------------------------------------------------------------
219 -- | O(segs). Yield a `UPSSegd` that describes each segment of a `UPVSegd`
220 -- individually.
221 --
222 -- * By doing this we lose information about virtual segments corresponding
223 -- to the same physical segments.
224 --
225 -- * This operation is used in concatPR as the first step in eliminating
226 -- segmentation from a nested array.
227 --
228 demoteToUPSSegd :: UPVSegd -> UPSSegd
229 demoteToUPSSegd (UPVSegd True _vsegids upssegd)
230 = upssegd
231
232 demoteToUPSSegd (UPVSegd False vsegids upssegd)
233 = let starts' = bpermuteUP (UPSSegd.takeStarts upssegd) vsegids
234 sources' = bpermuteUP (UPSSegd.takeSources upssegd) vsegids
235 lengths' = bpermuteUP (UPSSegd.takeLengths upssegd) vsegids
236 upsegd' = UPSegd.fromLengths lengths'
237 in UPSSegd.mkUPSSegd starts' sources' upsegd'
238 {-# NOINLINE demoteToUPSSegd #-}
239 -- NOINLINE because it's complicated and won't fuse with anything.
240 -- In core we want to see when VSegds are being demoted.
241
242
243 -- | O(segs). Given an virtual segment descriptor, produce a `UPSegd` that
244 -- that describes the entire array.
245 --
246 -- WARNING:
247 -- Trying to take the `UPSegd` of a nested array that has been constructed with
248 -- replication can cause index overflow. This is because the virtual size of
249 -- the corresponding flat data can be larger than physical memory.
250 --
251 -- You should only apply this function to a nested array when you're about
252 -- about to construct something with the same size as the corresponding
253 -- flat array. In this case the index overflow doesn't matter too much
254 -- because the program would OOM anyway.
255 --
256 unsafeDemoteToUPSegd :: UPVSegd -> UPSegd
257 unsafeDemoteToUPSegd (UPVSegd _ vsegids upssegd)
258 = UPSegd.fromLengths
259 $ bpermuteUP (UPSSegd.takeLengths upssegd) vsegids
260 {-# NOINLINE unsafeDemoteToUPSegd #-}
261 -- NOINLINE because it's complicated and won't fuse with anything.
262 -- In core we want to see when VSegds are being demoted.
263
264
265 -- Operators ------------------------------------------------------------------
266 -- | Update the virtual segment ids of a UPVSegd and force out unreachable
267 -- physical segments from the contained UPSSegd.
268 --
269 -- * TODO: make this parallel. It runs the sequential 'cull' then reconstructs
270 -- the UPSSegd.
271 --
272 updateVSegs :: (Vector Int -> Vector Int) -> UPVSegd -> UPVSegd
273 updateVSegs f (UPVSegd _ vsegids upssegd)
274 = let (vsegids', ussegd')
275 = USSegd.cullOnVSegids (f vsegids)
276 $ UPSSegd.takeUSSegd upssegd
277
278 in UPVSegd False vsegids' (UPSSegd.fromUSSegd ussegd')
279 {-# NOINLINE updateVSegs #-}
280
281
282 -- Append ---------------------------------------------------------------------
283 -- | Produce a segment descriptor that describes the result of appending two arrays.
284 --
285 -- * TODO: make this parallel.
286 --
287 appendWith
288 :: UPVSegd -> Int -- ^ uvsegd of array, and number of physical data arrays
289 -> UPVSegd -> Int -- ^ uvsegd of array, and number of physical data arrays
290 -> UPVSegd
291
292 appendWith
293 (UPVSegd _ vsegids1 upssegd1) pdatas1
294 (UPVSegd _ vsegids2 upssegd2) pdatas2
295
296 = let -- vsegids releative to appended psegs
297 vsegids1' = vsegids1
298 vsegids2' = V.map (+ UPSSegd.length upssegd1) vsegids2
299
300 -- append the vsegids
301 vsegids' = vsegids1' V.++ vsegids2'
302
303 -- All data from the source arrays goes into the result
304 upssegd' = UPSSegd.appendWith
305 upssegd1 pdatas1
306 upssegd2 pdatas2
307
308 in UPVSegd False vsegids' upssegd'
309 {-# NOINLINE appendWith #-}
310
311
312 -- Combine --------------------------------------------------------------------
313 -- | Combine two virtual segment descriptors.
314 --
315 -- * TODO: make this parallel.
316 --
317 combine2
318 :: UPSel2
319 -> UPVSegd -> Int -- ^ uvsegd of array, and number of physical data arrays
320 -> UPVSegd -> Int -- ^ uvsegd of array, and number of physical data arrays
321 -> UPVSegd
322
323 combine2
324 upsel2
325 (UPVSegd _ vsegids1 upssegd1) pdatas1
326 (UPVSegd _ vsegids2 upssegd2) pdatas2
327
328 = let -- vsegids relative to combined psegs
329 vsegids1' = vsegids1
330 vsegids2' = V.map (+ (V.length vsegids1)) vsegids2
331
332 -- combine the vsegids
333 vsegids' = V.combine2ByTag (UPSel.tagsUPSel2 upsel2)
334 vsegids1' vsegids2'
335
336 -- All data from the source arrays goes into the result
337 upssegd' = UPSSegd.appendWith
338 upssegd1 pdatas1
339 upssegd2 pdatas2
340
341 in UPVSegd False vsegids' upssegd'
342 {-# NOINLINE combine2 #-}
343