dph-prim-par: Add Justifications to distributed array functions
[packages/dph.git] / dph-prim-par / Data / Array / Parallel / Unlifted / Parallel / Enum.hs
1 {-# LANGUAGE CPP #-}
2 #include "fusion-phases.h"
3
4 -- | Enum-related parallel operations on unlifted arrays
5 module Data.Array.Parallel.Unlifted.Parallel.Enum
6 ( enumFromToUP
7 , enumFromThenToUP
8 , enumFromStepLenUP
9 , enumFromStepLenEachUP)
10 where
11 import Data.Array.Parallel.Unlifted.Sequential.Vector as Seq
12 import Data.Array.Parallel.Unlifted.Distributed
13 import Data.Array.Parallel.Unlifted.Distributed.What
14 import Data.Array.Parallel.Unlifted.Parallel.Combinators (mapUP)
15 import GHC.Base (divInt)
16
17
18 delay_inline :: a -> a
19 delay_inline x = x
20 {-# INLINE [0] delay_inline #-}
21
22
23 enumFromToUP :: (Unbox a, Enum a) => a -> a -> Vector a
24 enumFromToUP start end
25 = mapUP toEnum (enumFromStepLenUP start' 1 len)
26 where start' = fromEnum start
27 end' = fromEnum end
28 len = delay_inline max (end' - start' + 1) 0
29 {-# INLINE_UP enumFromToUP #-}
30
31
32 enumFromThenToUP :: (Unbox a, Enum a) => a -> a -> a -> Vector a
33 enumFromThenToUP start next end
34 = mapUP toEnum (enumFromStepLenUP start' delta len)
35 where start' = fromEnum start
36 next' = fromEnum next
37 end' = fromEnum end
38 delta = next' - start'
39
40 -- distance between start' and end' expressed in deltas
41 dist = (end' - start' + delta) `divInt` delta
42 len = max dist 0
43 {-# INLINE_UP enumFromThenToUP #-}
44
45
46 enumFromStepLenUP :: Int -> Int -> Int -> Vector Int
47 enumFromStepLenUP start delta len =
48 joinD theGang balanced
49 (mapD (What "enumFromStepLenUP/gen") theGang gen
50 (splitLenIdxD theGang len))
51 where
52 gen (n,i) = Seq.enumFromStepLen (i * delta + start) delta n
53 {-# INLINE_UP enumFromStepLenUP #-}
54
55
56 enumFromStepLenEachUP
57 :: Int -> Vector Int -> Vector Int -> Vector Int -> Vector Int
58 enumFromStepLenEachUP _n starts steps lens
59 = joinD theGang unbalanced
60 $ mapD (What "enumFromStepLenEachUP/enum") theGang enum
61 $ splitD theGang unbalanced (Seq.zip (Seq.zip starts steps) lens)
62 where
63 enum ps = let (qs, llens) = Seq.unzip ps
64 (lstarts, lsteps) = Seq.unzip qs
65 in Seq.enumFromStepLenEach (Seq.sum llens) lstarts lsteps llens
66 {-# INLINE_UP enumFromStepLenEachUP #-}
67