b7c44f98780ef6214404bf354d4e8ea16daf1849
[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 fUpdate (UPVSegd _ vsegids upssegd)
274 = let (vsegids', ussegd')
275 = USSegd.cullOnVSegids (fUpdate vsegids)
276 $ UPSSegd.takeUSSegd upssegd
277
278 in UPVSegd False vsegids' (UPSSegd.fromUSSegd ussegd')
279 {-# INLINE_UP updateVSegs #-}
280 -- INLINE_UP because we want to inline the parameter function fUpdate.
281
282
283 -- Append ---------------------------------------------------------------------
284 -- | Produce a segment descriptor that describes the result of appending two arrays.
285 --
286 -- * TODO: make this parallel.
287 --
288 appendWith
289 :: UPVSegd -> Int -- ^ uvsegd of array, and number of physical data arrays
290 -> UPVSegd -> Int -- ^ uvsegd of array, and number of physical data arrays
291 -> UPVSegd
292
293 appendWith
294 (UPVSegd _ vsegids1 upssegd1) pdatas1
295 (UPVSegd _ vsegids2 upssegd2) pdatas2
296
297 = let -- vsegids releative to appended psegs
298 vsegids1' = vsegids1
299 vsegids2' = V.map (+ UPSSegd.length upssegd1) vsegids2
300
301 -- append the vsegids
302 vsegids' = vsegids1' V.++ vsegids2'
303
304 -- All data from the source arrays goes into the result
305 upssegd' = UPSSegd.appendWith
306 upssegd1 pdatas1
307 upssegd2 pdatas2
308
309 in UPVSegd False vsegids' upssegd'
310 {-# NOINLINE appendWith #-}
311
312
313 -- Combine --------------------------------------------------------------------
314 -- | Combine two virtual segment descriptors.
315 --
316 -- * TODO: make this parallel.
317 --
318 combine2
319 :: UPSel2
320 -> UPVSegd -> Int -- ^ uvsegd of array, and number of physical data arrays
321 -> UPVSegd -> Int -- ^ uvsegd of array, and number of physical data arrays
322 -> UPVSegd
323
324 combine2
325 upsel2
326 (UPVSegd _ vsegids1 upssegd1) pdatas1
327 (UPVSegd _ vsegids2 upssegd2) pdatas2
328
329 = let -- vsegids relative to combined psegs
330 vsegids1' = vsegids1
331 vsegids2' = V.map (+ (V.length vsegids1)) vsegids2
332
333 -- combine the vsegids
334 vsegids' = V.combine2ByTag (UPSel.tagsUPSel2 upsel2)
335 vsegids1' vsegids2'
336
337 -- All data from the source arrays goes into the result
338 upssegd' = UPSSegd.appendWith
339 upssegd1 pdatas1
340 upssegd2 pdatas2
341
342 in UPVSegd False vsegids' upssegd'
343 {-# NOINLINE combine2 #-}
344