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
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
|
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
module Data.Array.Nested.Internal.Arith where
import Control.Monad (forM, guard)
import qualified Data.Array.Internal as OI
import qualified Data.Array.Internal.RankedG as RG
import qualified Data.Array.Internal.RankedS as RS
import Data.Bits
import Data.Int
import Data.List (sort)
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM
import Foreign.C.Types
import Foreign.Ptr
import Foreign.Storable (Storable)
import GHC.TypeLits
import Language.Haskell.TH
import System.IO.Unsafe
import Data.Array.Nested.Internal.Arith.Foreign
import Data.Array.Nested.Internal.Arith.Lists
liftVEltwise1 :: Storable a
=> SNat n
-> (VS.Vector a -> VS.Vector a)
-> RS.Array n a -> RS.Array n a
liftVEltwise1 SNat f arr@(RS.A (RG.A sh (OI.T strides offset vec)))
| Just prefixSz <- stridesDense sh strides =
let vec' = f (VS.slice offset prefixSz vec)
in RS.A (RG.A sh (OI.T strides 0 vec'))
| otherwise = RS.fromVector sh (f (RS.toVector arr))
liftVEltwise2 :: Storable a
=> SNat n
-> (Either a (VS.Vector a) -> Either a (VS.Vector a) -> VS.Vector a)
-> RS.Array n a -> RS.Array n a -> RS.Array n a
liftVEltwise2 SNat f
arr1@(RS.A (RG.A sh1 (OI.T strides1 offset1 vec1)))
arr2@(RS.A (RG.A sh2 (OI.T strides2 offset2 vec2)))
| sh1 /= sh2 = error $ "liftVEltwise2: shapes unequal: " ++ show sh1 ++ " vs " ++ show sh2
| product sh1 == 0 = arr1 -- if the arrays are empty, just return one of the empty inputs
| otherwise = case (stridesDense sh1 strides1, stridesDense sh2 strides2) of
(Just 1, Just 1) -> -- both are a (potentially replicated) scalar; just apply f to the scalars
let vec' = f (Left (vec1 VS.! offset1)) (Left (vec2 VS.! offset2))
in RS.A (RG.A sh1 (OI.T strides1 0 vec'))
(Just 1, Just n) -> -- scalar * dense
RS.fromVector sh1 (f (Left (vec1 VS.! offset1)) (Right (VS.slice offset2 n vec2)))
(Just n, Just 1) -> -- dense * scalar
RS.fromVector sh1 (f (Right (VS.slice offset1 n vec1)) (Left (vec2 VS.! offset2)))
(_, _) -> -- fallback case
RS.fromVector sh1 (f (Right (RS.toVector arr1)) (Right (RS.toVector arr2)))
-- | Given the shape vector and the stride vector, return whether this vector
-- of strides uses a dense prefix of its backing array. If so, the number of
-- elements in this prefix is returned.
-- This excludes any offset.
stridesDense :: [Int] -> [Int] -> Maybe Int
stridesDense sh _ | any (<= 0) sh = Just 0
stridesDense sh str =
-- sort dimensions on their stride, ascending, dropping any zero strides
case dropWhile ((== 0) . fst) (sort (zip str sh)) of
[] -> Just 1
(1, n) : (unzip -> (str', sh')) -> checkCover n sh' str'
_ -> Nothing -- if the smallest stride is not 1, it will never be dense
where
-- Given size of currently densely covered region at beginning of the
-- array, the remaining shape vector and the corresponding remaining stride
-- vector, return whether this all together covers a dense prefix of the
-- array. If it does, return the number of elements in this prefix.
checkCover :: Int -> [Int] -> [Int] -> Maybe Int
checkCover block [] [] = Just block
checkCover block (n : sh') (s : str') = guard (s <= block) >> checkCover (max block (n * s)) sh' str'
checkCover _ _ _ = error "Orthotope array's shape vector and stride vector have different lengths"
{-# NOINLINE vectorOp1 #-}
vectorOp1 :: forall a b. Storable a
=> (Ptr a -> Ptr b)
-> (Int64 -> Ptr b -> Ptr b -> IO ())
-> VS.Vector a -> VS.Vector a
vectorOp1 ptrconv f v = unsafePerformIO $ do
outv <- VSM.unsafeNew (VS.length v)
VSM.unsafeWith outv $ \poutv ->
VS.unsafeWith v $ \pv ->
f (fromIntegral (VS.length v)) (ptrconv poutv) (ptrconv pv)
VS.unsafeFreeze outv
-- | If two vectors are given, assumes that they have the same length.
{-# NOINLINE vectorOp2 #-}
vectorOp2 :: forall a b. Storable a
=> (a -> b)
-> (Ptr a -> Ptr b)
-> (a -> a -> a)
-> (Int64 -> Ptr b -> b -> Ptr b -> IO ()) -- sv
-> (Int64 -> Ptr b -> Ptr b -> b -> IO ()) -- vs
-> (Int64 -> Ptr b -> Ptr b -> Ptr b -> IO ()) -- vv
-> Either a (VS.Vector a) -> Either a (VS.Vector a) -> VS.Vector a
vectorOp2 valconv ptrconv fss fsv fvs fvv = \cases
(Left x) (Left y) -> VS.singleton (fss x y)
(Left x) (Right vy) ->
unsafePerformIO $ do
outv <- VSM.unsafeNew (VS.length vy)
VSM.unsafeWith outv $ \poutv ->
VS.unsafeWith vy $ \pvy ->
fsv (fromIntegral (VS.length vy)) (ptrconv poutv) (valconv x) (ptrconv pvy)
VS.unsafeFreeze outv
(Right vx) (Left y) ->
unsafePerformIO $ do
outv <- VSM.unsafeNew (VS.length vx)
VSM.unsafeWith outv $ \poutv ->
VS.unsafeWith vx $ \pvx ->
fvs (fromIntegral (VS.length vx)) (ptrconv poutv) (ptrconv pvx) (valconv y)
VS.unsafeFreeze outv
(Right vx) (Right vy)
| VS.length vx == VS.length vy ->
unsafePerformIO $ do
outv <- VSM.unsafeNew (VS.length vx)
VSM.unsafeWith outv $ \poutv ->
VS.unsafeWith vx $ \pvx ->
VS.unsafeWith vy $ \pvy ->
fvv (fromIntegral (VS.length vx)) (ptrconv poutv) (ptrconv pvx) (ptrconv pvy)
VS.unsafeFreeze outv
| otherwise -> error $ "vectorOp: unequal lengths: " ++ show (VS.length vx) ++ " /= " ++ show (VS.length vy)
-- TODO: test all the weird cases of this function
-- | Reduce along the inner dimension
{-# NOINLINE vectorRedInnerOp #-}
vectorRedInnerOp :: forall a b n. (Num a, Storable a)
=> SNat n
-> (a -> b)
-> (Ptr a -> Ptr b)
-> (Int64 -> Ptr b -> b -> Ptr b -> IO ()) -- ^ scale by constant
-> (Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr b -> Ptr b -> IO ()) -- ^ reduction kernel
-> RS.Array (n + 1) a -> RS.Array n a
vectorRedInnerOp sn@SNat valconv ptrconv fscale fred (RS.A (RG.A sh (OI.T strides offset vec)))
| null sh = error "unreachable"
| any (<= 0) sh = RS.A (RG.A (init sh) (OI.T (map (const 0) (init strides)) 0 VS.empty))
-- now the input array is nonempty
| last sh == 1 = RS.A (RG.A (init sh) (OI.T (init strides) offset vec))
| last strides == 0 =
liftVEltwise1 sn
(vectorOp1 id (\n pout px -> fscale n (ptrconv pout) (valconv (fromIntegral (last sh))) (ptrconv px)))
(RS.A (RG.A (init sh) (OI.T (init strides) offset vec)))
-- now there is useful work along the inner dimension
| otherwise =
let -- filter out zero-stride dimensions; the reduction kernel need not concern itself with those
(shF, stridesF) = unzip $ filter ((/= 0) . snd) (zip sh strides)
ndimsF = length shF
in unsafePerformIO $ do
outv <- VSM.unsafeNew (product (init shF))
VSM.unsafeWith outv $ \poutv ->
VS.unsafeWith (VS.fromListN ndimsF (map fromIntegral shF)) $ \pshF ->
VS.unsafeWith (VS.fromListN ndimsF (map fromIntegral stridesF)) $ \pstridesF ->
VS.unsafeWith (VS.slice offset (VS.length vec - offset) vec) $ \pvec ->
fred (fromIntegral ndimsF) pshF pstridesF (ptrconv poutv) (ptrconv pvec)
RS.fromVector (init sh) <$> VS.unsafeFreeze outv
flipOp :: (Int64 -> Ptr a -> a -> Ptr a -> IO ())
-> Int64 -> Ptr a -> Ptr a -> a -> IO ()
flipOp f n out v s = f n out s v
class NumElt a where
numEltAdd :: SNat n -> RS.Array n a -> RS.Array n a -> RS.Array n a
numEltSub :: SNat n -> RS.Array n a -> RS.Array n a -> RS.Array n a
numEltMul :: SNat n -> RS.Array n a -> RS.Array n a -> RS.Array n a
numEltNeg :: SNat n -> RS.Array n a -> RS.Array n a
numEltAbs :: SNat n -> RS.Array n a -> RS.Array n a
numEltSignum :: SNat n -> RS.Array n a -> RS.Array n a
numEltSum1Inner :: SNat n -> RS.Array (n + 1) a -> RS.Array n a
numEltProduct1Inner :: SNat n -> RS.Array (n + 1) a -> RS.Array n a
$(fmap concat . forM typesList $ \arithtype -> do
let ttyp = conT (atType arithtype)
fmap concat . forM binopsList $ \arithop -> do
let name = mkName (aboName arithop ++ "Vector" ++ nameBase (atType arithtype))
cnamebase = "c_" ++ aboName arithop ++ "_" ++ atCName arithtype
c_ss = varE (aboScalFun arithop arithtype)
c_sv = varE $ mkName (cnamebase ++ "_sv")
c_vs | aboComm arithop == NonComm = varE $ mkName (cnamebase ++ "_vs")
| otherwise = [| flipOp $c_sv |]
c_vv = varE $ mkName (cnamebase ++ "_vv")
sequence [SigD name <$>
[t| forall n. SNat n -> RS.Array n $ttyp -> RS.Array n $ttyp -> RS.Array n $ttyp |]
,do body <- [| \sn -> liftVEltwise2 sn (vectorOp2 id id $c_ss $c_sv $c_vs $c_vv) |]
return $ FunD name [Clause [] (NormalB body) []]])
$(fmap concat . forM typesList $ \arithtype -> do
let ttyp = conT (atType arithtype)
fmap concat . forM unopsList $ \arithop -> do
let name = mkName (auoName arithop ++ "Vector" ++ nameBase (atType arithtype))
c_op = varE $ mkName ("c_" ++ auoName arithop ++ "_" ++ atCName arithtype)
sequence [SigD name <$>
[t| forall n. SNat n -> RS.Array n $ttyp -> RS.Array n $ttyp |]
,do body <- [| \sn -> liftVEltwise1 sn (vectorOp1 id $c_op) |]
return $ FunD name [Clause [] (NormalB body) []]])
$(fmap concat . forM typesList $ \arithtype -> do
let ttyp = conT (atType arithtype)
fmap concat . forM redopsList $ \redop -> do
let name = mkName (aroName redop ++ "Vector" ++ nameBase (atType arithtype))
c_op = varE $ mkName ("c_" ++ aroName redop ++ "_" ++ atCName arithtype)
c_scale_op = varE $ mkName ("c_mul_" ++ atCName arithtype ++ "_sv")
sequence [SigD name <$>
[t| forall n. SNat n -> RS.Array (n + 1) $ttyp -> RS.Array n $ttyp |]
,do body <- [| \sn -> vectorRedInnerOp sn id id $c_scale_op $c_op |]
return $ FunD name [Clause [] (NormalB body) []]])
-- This branch is ostensibly a runtime branch, but will (hopefully) be
-- constant-folded away by GHC.
intWidBranch1 :: forall i n. (FiniteBits i, Storable i)
=> (Int64 -> Ptr Int32 -> Ptr Int32 -> IO ())
-> (Int64 -> Ptr Int64 -> Ptr Int64 -> IO ())
-> (SNat n -> RS.Array n i -> RS.Array n i)
intWidBranch1 f32 f64 sn
| finiteBitSize (undefined :: i) == 32 = liftVEltwise1 sn (vectorOp1 @i @Int32 castPtr f32)
| finiteBitSize (undefined :: i) == 64 = liftVEltwise1 sn (vectorOp1 @i @Int64 castPtr f64)
| otherwise = error "Unsupported Int width"
intWidBranch2 :: forall i n. (FiniteBits i, Storable i, Integral i)
=> (i -> i -> i) -- ss
-- int32
-> (Int64 -> Ptr Int32 -> Int32 -> Ptr Int32 -> IO ()) -- sv
-> (Int64 -> Ptr Int32 -> Ptr Int32 -> Int32 -> IO ()) -- vs
-> (Int64 -> Ptr Int32 -> Ptr Int32 -> Ptr Int32 -> IO ()) -- vv
-- int64
-> (Int64 -> Ptr Int64 -> Int64 -> Ptr Int64 -> IO ()) -- sv
-> (Int64 -> Ptr Int64 -> Ptr Int64 -> Int64 -> IO ()) -- vs
-> (Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int64 -> IO ()) -- vv
-> (SNat n -> RS.Array n i -> RS.Array n i -> RS.Array n i)
intWidBranch2 ss sv32 vs32 vv32 sv64 vs64 vv64 sn
| finiteBitSize (undefined :: i) == 32 = liftVEltwise2 sn (vectorOp2 @i @Int32 fromIntegral castPtr ss sv32 vs32 vv32)
| finiteBitSize (undefined :: i) == 64 = liftVEltwise2 sn (vectorOp2 @i @Int64 fromIntegral castPtr ss sv64 vs64 vv64)
| otherwise = error "Unsupported Int width"
intWidBranchRed :: forall i n. (FiniteBits i, Storable i, Integral i)
=> -- int32
(Int64 -> Ptr Int32 -> Int32 -> Ptr Int32 -> IO ()) -- ^ scale by constant
-> (Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int32 -> Ptr Int32 -> IO ()) -- ^ reduction kernel
-- int64
-> (Int64 -> Ptr Int64 -> Int64 -> Ptr Int64 -> IO ()) -- ^ scale by constant
-> (Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int64 -> IO ()) -- ^ reduction kernel
-> (SNat n -> RS.Array (n + 1) i -> RS.Array n i)
intWidBranchRed fsc32 fred32 fsc64 fred64 sn
| finiteBitSize (undefined :: i) == 32 = vectorRedInnerOp @i @Int32 sn fromIntegral castPtr fsc32 fred32
| finiteBitSize (undefined :: i) == 64 = vectorRedInnerOp @i @Int64 sn fromIntegral castPtr fsc64 fred64
| otherwise = error "Unsupported Int width"
instance NumElt Int32 where
numEltAdd = addVectorInt32
numEltSub = subVectorInt32
numEltMul = mulVectorInt32
numEltNeg = negVectorInt32
numEltAbs = absVectorInt32
numEltSignum = signumVectorInt32
numEltSum1Inner = sum1VectorInt32
numEltProduct1Inner = product1VectorInt32
instance NumElt Int64 where
numEltAdd = addVectorInt64
numEltSub = subVectorInt64
numEltMul = mulVectorInt64
numEltNeg = negVectorInt64
numEltAbs = absVectorInt64
numEltSignum = signumVectorInt64
numEltSum1Inner = sum1VectorInt64
numEltProduct1Inner = product1VectorInt64
instance NumElt Float where
numEltAdd = addVectorFloat
numEltSub = subVectorFloat
numEltMul = mulVectorFloat
numEltNeg = negVectorFloat
numEltAbs = absVectorFloat
numEltSignum = signumVectorFloat
numEltSum1Inner = sum1VectorFloat
numEltProduct1Inner = product1VectorFloat
instance NumElt Double where
numEltAdd = addVectorDouble
numEltSub = subVectorDouble
numEltMul = mulVectorDouble
numEltNeg = negVectorDouble
numEltAbs = absVectorDouble
numEltSignum = signumVectorDouble
numEltSum1Inner = sum1VectorDouble
numEltProduct1Inner = product1VectorDouble
instance NumElt Int where
numEltAdd = intWidBranch2 @Int (+) c_add_i32_sv (flipOp c_add_i32_sv) c_add_i32_vv c_add_i64_sv (flipOp c_add_i64_sv) c_add_i64_vv
numEltSub = intWidBranch2 @Int (-) c_sub_i32_sv (flipOp c_sub_i32_sv) c_sub_i32_vv c_sub_i64_sv (flipOp c_sub_i64_sv) c_sub_i64_vv
numEltMul = intWidBranch2 @Int (*) c_mul_i32_sv (flipOp c_mul_i32_sv) c_mul_i32_vv c_mul_i64_sv (flipOp c_mul_i64_sv) c_mul_i64_vv
numEltNeg = intWidBranch1 @Int c_neg_i32 c_neg_i64
numEltAbs = intWidBranch1 @Int c_abs_i32 c_abs_i64
numEltSignum = intWidBranch1 @Int c_signum_i32 c_signum_i64
numEltSum1Inner = intWidBranchRed @Int c_mul_i32_sv c_sum1_i32 c_mul_i64_sv c_sum1_i64
numEltProduct1Inner = intWidBranchRed @Int c_mul_i32_sv c_product1_i32 c_mul_i64_sv c_product1_i64
instance NumElt CInt where
numEltAdd = intWidBranch2 @CInt (+) c_add_i32_sv (flipOp c_add_i32_sv) c_add_i32_vv c_add_i64_sv (flipOp c_add_i64_sv) c_add_i64_vv
numEltSub = intWidBranch2 @CInt (-) c_sub_i32_sv (flipOp c_sub_i32_sv) c_sub_i32_vv c_sub_i64_sv (flipOp c_sub_i64_sv) c_sub_i64_vv
numEltMul = intWidBranch2 @CInt (*) c_mul_i32_sv (flipOp c_mul_i32_sv) c_mul_i32_vv c_mul_i64_sv (flipOp c_mul_i64_sv) c_mul_i64_vv
numEltNeg = intWidBranch1 @CInt c_neg_i32 c_neg_i64
numEltAbs = intWidBranch1 @CInt c_abs_i32 c_abs_i64
numEltSignum = intWidBranch1 @CInt c_signum_i32 c_signum_i64
numEltSum1Inner = intWidBranchRed @CInt c_mul_i32_sv c_sum1_i32 c_mul_i64_sv c_sum1_i64
numEltProduct1Inner = intWidBranchRed @CInt c_mul_i32_sv c_product1_i32 c_mul_i64_sv c_product1_i64
|
