1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
|
{-# LANGUAGE FunctionalDependencies, UndecidableInstances #-}
module Data.Array.Repa.Operators.Mapping
( -- * Generic maps
map
, zipWith
, (+^), (-^), (*^), (/^)
-- * Structured maps
, Structured(..))
where
import Data.Array.Repa.Shape
import Data.Array.Repa.Base
import Data.Array.Repa.Repr.ByteString
import Data.Array.Repa.Repr.Cursored
import Data.Array.Repa.Repr.Delayed
import Data.Array.Repa.Repr.ForeignPtr
import Data.Array.Repa.Repr.HintSmall
import Data.Array.Repa.Repr.HintInterleave
import Data.Array.Repa.Repr.Partitioned
import Data.Array.Repa.Repr.Unboxed
import Data.Array.Repa.Repr.Undefined
import Prelude hiding (map, zipWith)
import Foreign.Storable
import Data.Word
-- | Apply a worker function to each element of an array,
-- yielding a new array with the same extent.
--
map :: (Shape sh, Source r a)
=> (a -> b) -> Array r sh a -> Array D sh b
map f arr
= case delay arr of
ADelayed sh g -> ADelayed sh (f . g)
{-# INLINE [3] map #-}
-- ZipWith --------------------------------------------------------------------
-- | Combine two arrays, element-wise, with a binary operator.
-- If the extent of the two array arguments differ,
-- then the resulting array's extent is their intersection.
--
zipWith :: (Shape sh, Source r1 a, Source r2 b)
=> (a -> b -> c)
-> Array r1 sh a -> Array r2 sh b
-> Array D sh c
zipWith f arr1 arr2
= let get ix = f (arr1 `unsafeIndex` ix) (arr2 `unsafeIndex` ix)
{-# INLINE get #-}
in fromFunction
(intersectDim (extent arr1) (extent arr2))
get
{-# INLINE [2] zipWith #-}
infixl 7 *^, /^
infixl 6 +^, -^
(+^) = zipWith (+)
{-# INLINE (+^) #-}
(-^) = zipWith (-)
{-# INLINE (-^) #-}
(*^) = zipWith (*)
{-# INLINE (*^) #-}
(/^) = zipWith (/)
{-# INLINE (/^) #-}
-- Structured -------------------------------------------------------------------
-- | Structured versions of @map@ and @zipWith@ that preserve the representation
-- of cursored and partitioned arrays.
--
-- For cursored (@C@) arrays, the cursoring of the source array is preserved.
--
-- For partitioned (@P@) arrays, the worker function is fused with each array
-- partition separately, instead of treating the whole array as a single
-- bulk object.
--
-- Preserving the cursored and\/or paritioned representation of an array
-- is will make follow-on computation more efficient than if the array was
-- converted to a vanilla Delayed (@D@) array as with plain `map` and `zipWith`.
--
-- If the source array is not cursored or partitioned then `smap` and
-- `szipWith` are identical to the plain functions.
--
class Structured r1 a b where
-- | The target result representation.
type TR r1
-- | Structured @map@.
smap :: Shape sh
=> (a -> b)
-> Array r1 sh a
-> Array (TR r1) sh b
-- | Structured @zipWith@.
-- If you have a cursored or partitioned source array then use that as
-- the third argument (corresponding to @r1@ here)
szipWith
:: (Shape sh, Source r c)
=> (c -> a -> b)
-> Array r sh c
-> Array r1 sh a
-> Array (TR r1) sh b
-- ByteString -------------------------
instance Structured B Word8 b where
type TR B = D
smap = map
szipWith = zipWith
-- Cursored ---------------------------
instance Structured C a b where
type TR C = C
smap f (ACursored sh makec shiftc loadc)
= ACursored sh makec shiftc (f . loadc)
{-# INLINE [3] smap #-}
szipWith f arr1 (ACursored sh makec shiftc loadc)
= let makec' ix = (ix, makec ix)
{-# INLINE makec' #-}
shiftc' off (ix, cur) = (addDim off ix, shiftc off cur)
{-# INLINE shiftc' #-}
load' (ix, cur) = f (arr1 `unsafeIndex` ix) (loadc cur)
{-# INLINE load' #-}
in ACursored
(intersectDim (extent arr1) sh)
makec' shiftc' load'
{-# INLINE [2] szipWith #-}
-- Delayed ----------------------------
instance Structured D a b where
type TR D = D
smap = map
szipWith = zipWith
-- ForeignPtr -------------------------
instance Storable a => Structured F a b where
type TR F = D
smap = map
szipWith = zipWith
-- Partitioned ------------------------
instance (Structured r1 a b
, Structured r2 a b)
=> Structured (P r1 r2) a b where
type TR (P r1 r2) = P (TR r1) (TR r2)
smap f (APart sh range arr1 arr2)
= APart sh range (smap f arr1) (smap f arr2)
{-# INLINE [3] smap #-}
szipWith f arr1 (APart sh range arr21 arr22)
= APart sh range (szipWith f arr1 arr21)
(szipWith f arr1 arr22)
{-# INLINE [2] szipWith #-}
-- Small ------------------------------
instance Structured r1 a b
=> Structured (S r1) a b where
type TR (S r1) = S (TR r1)
smap f (ASmall arr1)
= ASmall (smap f arr1)
{-# INLINE [3] smap #-}
szipWith f arr1 (ASmall arr2)
= ASmall (szipWith f arr1 arr2)
{-# INLINE [3] szipWith #-}
-- Interleaved ------------------------
instance Structured r1 a b
=> Structured (I r1) a b where
type TR (I r1) = I (TR r1)
smap f (AInterleave arr1)
= AInterleave (smap f arr1)
{-# INLINE [3] smap #-}
szipWith f arr1 (AInterleave arr2)
= AInterleave (szipWith f arr1 arr2)
{-# INLINE [3] szipWith #-}
-- Unboxed ----------------------------
instance Unbox a => Structured U a b where
type TR U = D
smap = map
szipWith = zipWith
-- Undefined --------------------------
instance Structured X a b where
type TR X = X
smap _ (AUndefined sh) = AUndefined sh
szipWith _ _ (AUndefined sh) = AUndefined sh
|
