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
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
|
{-# LANGUAGE CPP, MagicHash #-}
-- This is specialised for stencils up to 7x7.
-- Due to limitations in the GHC optimiser, using larger stencils doesn't
-- work, and will yield `error` at runtime. We can probably increase the
-- limit if required -- just ask.
--
-- The focus of the stencil is in the center of the 7x7 tile, which has
-- coordinates (0, 0). All coefficients in the stencil must fit in the tile,
-- so they can be given X,Y coordinates up to +/- 3 positions.
-- The stencil can be any shape, and need not be symmetric -- provided it
-- fits in the 7x7 tile.
--
module Data.Array.Repa.Stencil.Dim2
( -- * Stencil creation
makeStencil2,
#ifndef REPA_NO_TH
stencil2,
#endif
-- * Stencil operators
PC5, mapStencil2, forStencil2)
where
import Data.Array.Repa.Base
import Data.Array.Repa.Index
import Data.Array.Repa.Shape
import Data.Array.Repa.Repr.Delayed
import Data.Array.Repa.Repr.Cursored
import Data.Array.Repa.Repr.Partitioned
import Data.Array.Repa.Repr.HintSmall
import Data.Array.Repa.Repr.Undefined
import Data.Array.Repa.Stencil.Base
#ifndef REPA_NO_TH
import Data.Array.Repa.Stencil.Template
#endif
import Data.Array.Repa.Stencil.Partition
import GHC.Exts
-- | A index into the flat array.
-- Should be abstract outside the stencil modules.
data Cursor
= Cursor Int
type PC5 = P C (P (S D) (P (S D) (P (S D) (P (S D) X))))
-- Wrappers -------------------------------------------------------------------
-- | Like `mapStencil2` but with the parameters flipped.
forStencil2
:: Source r a
=> Boundary a
-> Array r DIM2 a
-> Stencil DIM2 a
-> Array PC5 DIM2 a
{-# INLINE forStencil2 #-}
forStencil2 boundary arr stencil
= mapStencil2 boundary stencil arr
-------------------------------------------------------------------------------
-- | Apply a stencil to every element of a 2D array.
mapStencil2
:: Source r a
=> Boundary a -- ^ How to handle the boundary of the array.
-> Stencil DIM2 a -- ^ Stencil to apply.
-> Array r DIM2 a -- ^ Array to apply stencil to.
-> Array PC5 DIM2 a
{-# INLINE mapStencil2 #-}
mapStencil2 boundary stencil@(StencilStatic sExtent _zero _load) arr
= let sh = extent arr
(_ :. aHeight :. aWidth) = sh
(_ :. sHeight :. sWidth) = sExtent
sHeight2 = sHeight `div` 2
sWidth2 = sWidth `div` 2
-- Partition the array into the internal and border regions.
![ Region inX inY inW inH
, Region westX westY westW westH
, Region eastX eastY eastW eastH
, Region northX northY northW northH
, Region southX southY southW southH ]
= partitionForStencil
(Size aWidth aHeight)
(Size sWidth sHeight)
(Offset sWidth2 sHeight2)
{-# INLINE inInternal #-}
inInternal (Z :. y :. x)
= x >= inX && x < (inX + inW)
&& y >= inY && y < (inY + inH)
{-# INLINE inBorder #-}
inBorder = not . inInternal
-- Cursor functions ----------------
{-# INLINE makec #-}
makec (Z :. y :. x)
= Cursor (x + y * aWidth)
{-# INLINE shiftc #-}
shiftc ix (Cursor off)
= Cursor
$ case ix of
Z :. y :. x -> off + y * aWidth + x
{-# INLINE arrInternal #-}
arrInternal = makeCursored (extent arr) makec shiftc getInner'
{-# INLINE getInner' #-}
getInner' cur = unsafeAppStencilCursor2 shiftc stencil arr cur
{-# INLINE arrBorder #-}
arrBorder = ASmall (fromFunction (extent arr) getBorder')
{-# INLINE getBorder' #-}
getBorder' ix
= case boundary of
BoundFixed c -> c
BoundConst c -> unsafeAppStencilCursor2_const addDim stencil c arr ix
BoundClamp -> unsafeAppStencilCursor2_clamp addDim stencil arr ix
in
-- internal region
APart sh (Range (Z :. inY :. inX) (Z :. inH :. inW) inInternal) arrInternal
-- border regions
$ APart sh (Range (Z :. westY :. westX) (Z :. westH :. westW) inBorder) arrBorder
$ APart sh (Range (Z :. eastY :. eastX) (Z :. eastH :. eastW) inBorder) arrBorder
$ APart sh (Range (Z :. northY :. northX) (Z :. northH :. northW) inBorder) arrBorder
$ APart sh (Range (Z :. southY :. southX) (Z :. southH :. southW) inBorder) arrBorder
$ AUndefined sh
unsafeAppStencilCursor2
:: Source r a
=> (DIM2 -> Cursor -> Cursor)
-> Stencil DIM2 a
-> Array r DIM2 a
-> Cursor
-> a
{-# INLINE unsafeAppStencilCursor2 #-}
unsafeAppStencilCursor2 shift
(StencilStatic sExtent zero loads)
arr cur0
| _ :. sHeight :. sWidth <- sExtent
, sHeight <= 7, sWidth <= 7
= let
-- Get data from the manifest array.
{-# INLINE getData #-}
getData (Cursor cur) = arr `unsafeLinearIndex` cur
-- Build a function to pass data from the array to our stencil.
{-# INLINE oload #-}
oload oy ox
= let !cur' = shift (Z :. oy :. ox) cur0
in loads (Z :. oy :. ox) (getData cur')
in template7x7 oload zero
| otherwise
= error $ unlines
[ "mapStencil2: Your stencil is too big for this method."
, " It must fit within a 7x7 tile to be compiled statically." ]
-- | Like above, but treat elements outside the array has having a constant value.
unsafeAppStencilCursor2_const
:: forall r a
. Source r a
=> (DIM2 -> DIM2 -> DIM2)
-> Stencil DIM2 a
-> a
-> Array r DIM2 a
-> DIM2
-> a
{-# INLINE unsafeAppStencilCursor2_const #-}
unsafeAppStencilCursor2_const shift
(StencilStatic sExtent zero loads)
fixed arr cur
| _ :. sHeight :. sWidth <- sExtent
, _ :. (I# aHeight) :. (I# aWidth) <- extent arr
, sHeight <= 7, sWidth <= 7
= let
-- Get data from the manifest array.
{-# INLINE getData #-}
getData :: DIM2 -> a
getData (Z :. (I# y) :. (I# x))
= getData' x y
{-# NOINLINE getData' #-}
getData' :: Int# -> Int# -> a
getData' !x !y
| 1# <- (x <# 0#) `orI#` (x >=# aWidth)
`orI#` (y <# 0#) `orI#` (y >=# aHeight)
= fixed
| otherwise
= arr `unsafeIndex` (Z :. (I# y) :. (I# x))
-- Build a function to pass data from the array to our stencil.
{-# INLINE oload #-}
oload oy ox
= let !cur' = shift (Z :. oy :. ox) cur
in loads (Z :. oy :. ox) (getData cur')
in template7x7 oload zero
| otherwise
= error $ unlines
[ "mapStencil2: Your stencil is too big for this method."
, " It must fit within a 7x7 tile to be compiled statically." ]
-- | Like above, but clamp out of bounds array values to the closest real value.
unsafeAppStencilCursor2_clamp
:: forall r a
. Source r a
=> (DIM2 -> DIM2 -> DIM2)
-> Stencil DIM2 a
-> Array r DIM2 a
-> DIM2
-> a
{-# INLINE unsafeAppStencilCursor2_clamp #-}
unsafeAppStencilCursor2_clamp shift
(StencilStatic sExtent zero loads)
arr cur
| _ :. sHeight :. sWidth <- sExtent
, _ :. (I# aHeight) :. (I# aWidth) <- extent arr
, sHeight <= 7, sWidth <= 7
= let
-- Get data from the manifest array.
{-# INLINE getData #-}
getData :: DIM2 -> a
getData (Z :. (I# y) :. (I# x))
= wrapLoadX x y
{-# NOINLINE wrapLoadX #-}
wrapLoadX :: Int# -> Int# -> a
wrapLoadX !x !y
| 1# <- x <# 0# = wrapLoadY 0# y
| 1# <- x >=# aWidth = wrapLoadY (aWidth -# 1#) y
| otherwise = wrapLoadY x y
{-# NOINLINE wrapLoadY #-}
wrapLoadY :: Int# -> Int# -> a
wrapLoadY !x !y
| 1# <- y <# 0# = loadXY x 0#
| 1# <- y >=# aHeight = loadXY x (aHeight -# 1#)
| otherwise = loadXY x y
{-# INLINE loadXY #-}
loadXY :: Int# -> Int# -> a
loadXY !x !y
= arr `unsafeIndex` (Z :. (I# y) :. (I# x))
-- Build a function to pass data from the array to our stencil.
{-# INLINE oload #-}
oload oy ox
= let !cur' = shift (Z :. oy :. ox) cur
in loads (Z :. oy :. ox) (getData cur')
in template7x7 oload zero
| otherwise
= error $ unlines
[ "mapStencil2: Your stencil is too big for this method."
, " It must fit within a 7x7 tile to be compiled statically." ]
-- | Data template for stencils up to 7x7.
template7x7
:: (Int -> Int -> a -> a)
-> a -> a
{-# INLINE template7x7 #-}
template7x7 f zero
= f (-3) (-3) $ f (-3) (-2) $ f (-3) (-1) $ f (-3) 0 $ f (-3) 1 $ f (-3) 2 $ f (-3) 3
$ f (-2) (-3) $ f (-2) (-2) $ f (-2) (-1) $ f (-2) 0 $ f (-2) 1 $ f (-2) 2 $ f (-2) 3
$ f (-1) (-3) $ f (-1) (-2) $ f (-1) (-1) $ f (-1) 0 $ f (-1) 1 $ f (-1) 2 $ f (-1) 3
$ f 0 (-3) $ f 0 (-2) $ f 0 (-1) $ f 0 0 $ f 0 1 $ f 0 2 $ f 0 3
$ f 1 (-3) $ f 1 (-2) $ f 1 (-1) $ f 1 0 $ f 1 1 $ f 1 2 $ f 1 3
$ f 2 (-3) $ f 2 (-2) $ f 2 (-1) $ f 2 0 $ f 2 1 $ f 2 2 $ f 2 3
$ f 3 (-3) $ f 3 (-2) $ f 3 (-1) $ f 3 0 $ f 3 1 $ f 3 2 $ f 3 3
$ zero
|
