{-# LANGUAGE DataKinds #-} {-# LANGUAGE ImportQualifiedPost #-} {-# LANGUAGE LambdaCase #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE TupleSections #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE ViewPatterns #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} module Data.Array.Mixed.Internal.Arith where import Control.Monad (forM, guard) import Data.Array.Internal qualified as OI import Data.Array.Internal.RankedG qualified as RG import Data.Array.Internal.RankedS qualified as RS import Data.Bifunctor (second) import Data.Bits import Data.Int import Data.List (sort) import Data.Vector.Storable qualified as VS import Data.Vector.Storable.Mutable qualified as VSM import Foreign.C.Types import Foreign.Marshal.Alloc (alloca) import Foreign.Ptr import Foreign.Storable (Storable(sizeOf), peek, poke) import GHC.TypeLits import GHC.TypeNats qualified as TypeNats import Language.Haskell.TH import System.IO.Unsafe import Data.Array.Mixed.Internal.Arith.Foreign import Data.Array.Mixed.Internal.Arith.Lists import Data.Array.Mixed.Types (fromSNat') -- TODO: need to sort strides for reduction-like functions so that the C inner-loop specialisation has some chance of working even after transposition -- TODO: test all the cases of this thing with various input strides liftVEltwise1 :: (Storable a, Storable b) => SNat n -> (VS.Vector a -> VS.Vector b) -> RS.Array n a -> RS.Array n b liftVEltwise1 SNat f arr@(RS.A (RG.A sh (OI.T strides offset vec))) | Just (blockOff, blockSz) <- stridesDense sh offset strides = let vec' = f (VS.slice blockOff blockSz vec) in RS.A (RG.A sh (OI.T strides (offset - blockOff) vec')) | otherwise = RS.fromVector sh (f (RS.toVector arr)) -- TODO: test all the cases of this thing with various input strides liftVEltwise2 :: (Storable a, Storable b, Storable c) => SNat n -> (Either a (VS.Vector a) -> Either b (VS.Vector b) -> VS.Vector c) -> RS.Array n a -> RS.Array n b -> RS.Array n c 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 = RS.A (RG.A sh1 (OI.T (0 <$ strides1) 0 VS.empty)) | otherwise = case (stridesDense sh1 offset1 strides1, stridesDense sh2 offset2 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 (blockOff, blockSz)) -> -- scalar * dense RS.A (RG.A sh1 (OI.T strides2 (offset2 - blockOff) (f (Left (vec1 VS.! offset1)) (Right (VS.slice blockOff blockSz vec2))))) (Just (blockOff, blockSz), Just (_, 1)) -> -- dense * scalar RS.A (RG.A sh1 (OI.T strides1 (offset1 - blockOff) (f (Right (VS.slice blockOff blockSz vec1)) (Left (vec2 VS.! offset2))))) (Just (blockOff1, blockSz1), Just (blockOff2, blockSz2)) | blockSz1 == blockSz2 -- not sure if this check is necessary, might be implied by the below , strides1 == strides2 -> -- dense * dense but the strides match RS.A (RG.A sh1 (OI.T strides1 (offset1 - blockOff1) (f (Right (VS.slice blockOff1 blockSz1 vec1)) (Right (VS.slice blockOff2 blockSz2 vec2))))) (_, _) -> -- fallback case RS.fromVector sh1 (f (Right (RS.toVector arr1)) (Right (RS.toVector arr2))) -- | Given shape vector, offset and stride vector, check whether this virtual -- vector uses a dense subarray of its backing array. If so, the first index -- and the number of elements in this subarray is returned. -- This excludes any offset. stridesDense :: [Int] -> Int -> [Int] -> Maybe (Int, Int) stridesDense sh offset _ | any (<= 0) sh = Just (offset, 0) stridesDense sh offsetNeg stridesNeg = -- First reverse all dimensions with negative stride, so that the first used -- value is at 'offset' and the rest is >= offset. let (offset, strides) = flipReverseds sh offsetNeg stridesNeg in -- sort dimensions on their stride, ascending, dropping any zero strides case filter ((/= 0) . fst) (sort (zip strides sh)) of [] -> Just (offset, 1) (1, n) : pairs -> (offset,) <$> checkCover n pairs _ -> 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 and the remaining (stride, size) pairs with all strides >=1, -- 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 ((s, n) : pairs) = guard (s <= block) >> checkCover (max block (n * s)) pairs -- Given shape, offset and strides, returns new (offset, strides) such that all strides are >=0 flipReverseds :: [Int] -> Int -> [Int] -> (Int, [Int]) flipReverseds [] off [] = (off, []) flipReverseds (n : sh') off (s : str') | s >= 0 = second (s :) (flipReverseds sh' off str') | otherwise = let off' = off + (n - 1) * s in second ((-s) :) (flipReverseds sh' off' str') flipReverseds _ _ _ = error "flipReverseds: invalid arguments" {-# 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 handling of negative strides -- | 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 b -> Ptr Int64 -> Ptr Int64 -> 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" | last sh <= 0 = RS.stretch (init sh) (RS.fromList (1 <$ init sh) [0]) | any (<= 0) (init sh) = RS.A (RG.A (init sh) (OI.T (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 -- replicated dimensions: dimensions with zero stride. The reduction -- kernel need not concern itself with those (and in fact has a -- precondition that there are no such dimensions in its input). replDims = map (== 0) strides -- filter out replicated dimensions (shF, stridesF) = unzip [(n, s) | (n, s, False) <- zip3 sh strides replDims] -- replace replicated dimensions with ones shOnes = zipWith (\n repl -> if repl then 1 else n) sh replDims ndimsF = length shF -- > 0, otherwise `last strides == 0` -- reversed dimensions: dimensions with negative stride. Reversal is -- irrelevant for a reduction, and indeed the kernel has a -- precondition that there are no such dimensions. revDims = map (< 0) stridesF stridesR = map abs stridesF offsetR = offset + sum (zipWith3 (\rev n s -> if rev then (n - 1) * s else 0) revDims shF stridesF) -- The *R values give an array with strides all > 0, hence the -- left-most element is at offsetR. in unsafePerformIO $ do outvR <- VSM.unsafeNew (product (init shF)) VSM.unsafeWith outvR $ \poutvR -> VS.unsafeWith (VS.fromListN ndimsF (map fromIntegral shF)) $ \pshF -> VS.unsafeWith (VS.fromListN ndimsF (map fromIntegral stridesR)) $ \pstridesR -> VS.unsafeWith (VS.slice offsetR (VS.length vec - offsetR) vec) $ \pvecR -> fred (fromIntegral ndimsF) (ptrconv poutvR) pshF pstridesR (ptrconv pvecR) TypeNats.withSomeSNat (fromIntegral (ndimsF - 1)) $ \(SNat :: SNat lenFm1) -> RS.stretch (init sh) -- replicate to original shape . RS.reshape (init shOnes) -- add 1-sized dimensions where the original was replicated . RS.rev (map fst (filter snd (zip [0..] revDims))) -- re-reverse the correct dimensions . RS.fromVector @_ @lenFm1 (init shF) -- the partially-reversed result array <$> VS.unsafeFreeze outvR -- TODO: test handling of negative strides -- | Reduce full array {-# NOINLINE vectorRedFullOp #-} vectorRedFullOp :: forall a b n. (Num a, Storable a) => SNat n -> (a -> Int -> a) -> (b -> a) -> (Ptr a -> Ptr b) -> (Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr b -> IO b) -- ^ reduction kernel -> RS.Array n a -> a vectorRedFullOp _ scaleval valbackconv ptrconv fred (RS.A (RG.A sh (OI.T strides offset vec))) | null sh = vec VS.! offset -- 0D array has one element | any (<= 0) sh = 0 -- now the input array is nonempty | all (== 0) strides = fromIntegral (product sh) * vec VS.! offset -- now there is at least one non-replicated dimension | otherwise = let -- replicated dimensions: dimensions with zero stride. The reduction -- kernel need not concern itself with those (and in fact has a -- precondition that there are no such dimensions in its input). replDims = map (== 0) strides -- filter out replicated dimensions (shF, stridesF) = unzip [(n, s) | (n, s, False) <- zip3 sh strides replDims] ndimsF = length shF -- > 0, otherwise `all (== 0) strides` -- we should scale up the output this many times to account for the replicated dimensions multiplier = product [n | (n, True) <- zip sh replDims] -- reversed dimensions: dimensions with negative stride. Reversal is -- irrelevant for a reduction, and indeed the kernel has a -- precondition that there are no such dimensions. revDims = map (< 0) stridesF stridesR = map abs stridesF offsetR = offset + sum (zipWith3 (\rev n s -> if rev then (n - 1) * s else 0) revDims shF stridesF) -- The *R values give an array with strides all > 0, hence the -- left-most element is at offsetR. in unsafePerformIO $ do VS.unsafeWith (VS.fromListN ndimsF (map fromIntegral shF)) $ \pshF -> VS.unsafeWith (VS.fromListN ndimsF (map fromIntegral stridesR)) $ \pstridesR -> VS.unsafeWith (VS.slice offsetR (VS.length vec - offsetR) vec) $ \pvecR -> (`scaleval` fromIntegral multiplier) . valbackconv <$> fred (fromIntegral ndimsF) pshF pstridesR (ptrconv pvecR) -- TODO: test this function -- | Find extremum (minindex ("argmin") or maxindex) in full array {-# NOINLINE vectorExtremumOp #-} vectorExtremumOp :: forall a b n. Storable a => (Ptr a -> Ptr b) -> (Ptr Int64 -> Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr b -> IO ()) -- ^ extremum kernel -> RS.Array n a -> [Int] -- result length: n vectorExtremumOp ptrconv fextrem (RS.A (RG.A sh (OI.T strides offset vec))) | null sh = [] | any (<= 0) sh = error "Extremum (minindex/maxindex): empty array" -- now the input array is nonempty | all (== 0) strides = 0 <$ sh -- now there is at least one non-replicated dimension | otherwise = let -- replicated dimensions: dimensions with zero stride. The extremum -- kernel need not concern itself with those (and in fact has a -- precondition that there are no such dimensions in its input). replDims = map (== 0) strides -- filter out replicated dimensions (shF, stridesF) = unzip [(n, s) | (n, s, False) <- zip3 sh strides replDims] ndimsF = length shF -- > 0, because not all strides were <=0 -- un-reverse reversed dimensions revDims = map (< 0) stridesF stridesR = map abs stridesF offsetR = offset + sum (zipWith3 (\rev n s -> if rev then (n - 1) * s else 0) revDims shF stridesF) -- function to insert zeros in replicated-out dimensions insertZeros :: [Bool] -> [Int] -> [Int] insertZeros [] idx = idx insertZeros (True : repls) idx = 0 : insertZeros repls idx insertZeros (False : repls) (i : idx) = i : insertZeros repls idx insertZeros (_:_) [] = error "unreachable" in unsafePerformIO $ do outvR <- VSM.unsafeNew (length shF) VSM.unsafeWith outvR $ \poutvR -> VS.unsafeWith (VS.fromListN ndimsF (map fromIntegral shF)) $ \pshF -> VS.unsafeWith (VS.fromListN ndimsF (map fromIntegral stridesR)) $ \pstridesR -> VS.unsafeWith (VS.slice offsetR (VS.length vec - offsetR) vec) $ \pvecR -> fextrem poutvR (fromIntegral ndimsF) pshF pstridesR (ptrconv pvecR) insertZeros replDims . zipWith3 (\rev n i -> if rev then n - 1 - i else i) revDims shF -- re-reverse the reversed dimensions . map (fromIntegral @Int64 @Int) . VS.toList <$> VS.unsafeFreeze outvR vectorDotprodInnerOp :: forall a b n. (Num a, Storable a) => SNat n -> (a -> b) -> (Ptr a -> Ptr b) -> (SNat n -> RS.Array n a -> RS.Array n a -> RS.Array n a) -- ^ elementwise multiplication -> (Int64 -> Ptr b -> b -> Ptr b -> IO ()) -- ^ scale by constant -> (Int64 -> Ptr b -> Ptr Int64 -> Ptr Int64 -> Ptr b -> IO ()) -- ^ reduction kernel -> (Int64 -> Ptr Int64 -> Ptr b -> Ptr Int64 -> Ptr b -> Ptr Int64 -> Ptr b -> IO ()) -- ^ dotprod kernel -> RS.Array (n + 1) a -> RS.Array (n + 1) a -> RS.Array n a vectorDotprodInnerOp sn@SNat valconv ptrconv fmul fscale fred fdotinner arr1@(RS.A (RG.A sh1 (OI.T strides1 offset1 vec1))) arr2@(RS.A (RG.A sh2 (OI.T strides2 offset2 vec2))) | null sh1 || null sh2 = error "unreachable" | sh1 /= sh2 = error $ "vectorDotprodInnerOp: shapes unequal: " ++ show sh1 ++ " vs " ++ show sh2 | last sh1 <= 0 = RS.stretch (init sh1) (RS.fromList (1 <$ init sh1) [0]) | any (<= 0) (init sh1) = RS.A (RG.A (init sh1) (OI.T (0 <$ init strides1) 0 VS.empty)) -- now the input arrays are nonempty | last sh1 == 1 = fmul sn (RS.reshape (init sh1) arr1) (RS.reshape (init sh1) arr2) | last strides1 == 0 = fmul sn (RS.A (RG.A (init sh1) (OI.T (init strides1) offset1 vec1))) (vectorRedInnerOp sn valconv ptrconv fscale fred arr2) | last strides2 == 0 = fmul sn (vectorRedInnerOp sn valconv ptrconv fscale fred arr1) (RS.A (RG.A (init sh2) (OI.T (init strides2) offset2 vec2))) -- now there is useful dotprod work along the inner dimension | otherwise = unsafePerformIO $ do let inrank = fromSNat' sn + 1 outv <- VSM.unsafeNew (product (init sh1)) VSM.unsafeWith outv $ \poutv -> VS.unsafeWith (VS.fromListN inrank (map fromIntegral sh1)) $ \psh -> VS.unsafeWith (VS.fromListN inrank (map fromIntegral strides1)) $ \pstrides1 -> VS.unsafeWith vec1 $ \pvec1 -> VS.unsafeWith (VS.fromListN inrank (map fromIntegral strides2)) $ \pstrides2 -> VS.unsafeWith vec2 $ \pvec2 -> fdotinner (fromIntegral @Int @Int64 inrank) psh (ptrconv poutv) pstrides1 (ptrconv pvec1 `plusPtr` (sizeOf (undefined :: a) * offset1)) pstrides2 (ptrconv pvec2 `plusPtr` (sizeOf (undefined :: a) * offset2)) RS.fromVector @_ @n (init sh1) <$> VS.unsafeFreeze outv {-# NOINLINE dotScalarVector #-} dotScalarVector :: forall a b. (Num a, Storable a) => Int -> (Ptr a -> Ptr b) -> (Int64 -> Ptr b -> Ptr Int64 -> Ptr Int64 -> Ptr b -> IO ()) -- ^ reduction kernel -> a -> VS.Vector a -> a dotScalarVector len ptrconv fred scalar vec = unsafePerformIO $ do alloca @a $ \pout -> do alloca @Int64 $ \pshape -> do poke pshape (fromIntegral @Int @Int64 len) alloca @Int64 $ \pstride -> do poke pstride 1 VS.unsafeWith vec $ \pvec -> fred 1 (ptrconv pout) pshape pstride (ptrconv pvec) res <- peek pout return (scalar * res) {-# NOINLINE dotVectorVector #-} dotVectorVector :: Storable a => Int -> (b -> a) -> (Ptr a -> Ptr b) -> (Int64 -> Ptr b -> Ptr b -> IO b) -- ^ dotprod kernel -> VS.Vector a -> VS.Vector a -> a dotVectorVector len valbackconv ptrconv fdot vec1 vec2 = unsafePerformIO $ do VS.unsafeWith vec1 $ \pvec1 -> VS.unsafeWith vec2 $ \pvec2 -> valbackconv <$> fdot (fromIntegral @Int @Int64 len) (ptrconv pvec1) (ptrconv pvec2) {-# NOINLINE dotVectorVectorStrided #-} dotVectorVectorStrided :: Storable a => Int -> (b -> a) -> (Ptr a -> Ptr b) -> (Int64 -> Int64 -> Int64 -> Ptr b -> Int64 -> Int64 -> Ptr b -> IO b) -- ^ dotprod kernel -> Int -> Int -> VS.Vector a -> Int -> Int -> VS.Vector a -> a dotVectorVectorStrided len valbackconv ptrconv fdot offset1 stride1 vec1 offset2 stride2 vec2 = unsafePerformIO $ do VS.unsafeWith vec1 $ \pvec1 -> VS.unsafeWith vec2 $ \pvec2 -> valbackconv <$> fdot (fromIntegral @Int @Int64 len) (fromIntegral offset1) (fromIntegral stride1) (ptrconv pvec1) (fromIntegral offset2) (fromIntegral stride2) (ptrconv pvec2) 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 $(fmap concat . forM typesList $ \arithtype -> do let ttyp = conT (atType arithtype) fmap concat . forM [minBound..maxBound] $ \arithop -> do let name = mkName (aboName arithop ++ "Vector" ++ nameBase (atType arithtype)) cnamebase = "c_binary_" ++ atCName arithtype c_ss = varE (aboNumOp arithop) c_sv = varE (mkName (cnamebase ++ "_sv")) `appE` litE (integerL (fromIntegral (aboEnum arithop))) c_vs = varE (mkName (cnamebase ++ "_vs")) `appE` litE (integerL (fromIntegral (aboEnum arithop))) c_vv = varE (mkName (cnamebase ++ "_vv")) `appE` litE (integerL (fromIntegral (aboEnum arithop))) 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 floatTypesList $ \arithtype -> do let ttyp = conT (atType arithtype) fmap concat . forM [minBound..maxBound] $ \arithop -> do let name = mkName (afboName arithop ++ "Vector" ++ nameBase (atType arithtype)) cnamebase = "c_fbinary_" ++ atCName arithtype c_ss = varE (afboNumOp arithop) c_sv = varE (mkName (cnamebase ++ "_sv")) `appE` litE (integerL (fromIntegral (afboEnum arithop))) c_vs = varE (mkName (cnamebase ++ "_vs")) `appE` litE (integerL (fromIntegral (afboEnum arithop))) c_vv = varE (mkName (cnamebase ++ "_vv")) `appE` litE (integerL (fromIntegral (afboEnum arithop))) 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 [minBound..maxBound] $ \arithop -> do let name = mkName (auoName arithop ++ "Vector" ++ nameBase (atType arithtype)) c_op = varE (mkName ("c_unary_" ++ atCName arithtype)) `appE` litE (integerL (fromIntegral (auoEnum arithop))) 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 floatTypesList $ \arithtype -> do let ttyp = conT (atType arithtype) fmap concat . forM [minBound..maxBound] $ \arithop -> do let name = mkName (afuoName arithop ++ "Vector" ++ nameBase (atType arithtype)) c_op = varE (mkName ("c_funary_" ++ atCName arithtype)) `appE` litE (integerL (fromIntegral (afuoEnum arithop))) 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) []]]) mulWithInt :: Num a => a -> Int -> a mulWithInt a i = a * fromIntegral i $(fmap concat . forM typesList $ \arithtype -> do let ttyp = conT (atType arithtype) fmap concat . forM [minBound..maxBound] $ \arithop -> do let scaleVar = case arithop of RO_SUM -> varE 'mulWithInt RO_PRODUCT -> varE '(^) let name1 = mkName (aroName arithop ++ "1Vector" ++ nameBase (atType arithtype)) namefull = mkName (aroName arithop ++ "FullVector" ++ nameBase (atType arithtype)) c_op1 = varE (mkName ("c_reduce1_" ++ atCName arithtype)) `appE` litE (integerL (fromIntegral (aroEnum arithop))) c_opfull = varE (mkName ("c_reducefull_" ++ atCName arithtype)) `appE` litE (integerL (fromIntegral (aroEnum arithop))) c_scale_op = varE (mkName ("c_binary_" ++ atCName arithtype ++ "_sv")) `appE` litE (integerL (fromIntegral (aboEnum BO_MUL))) sequence [SigD name1 <$> [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_op1 |] return $ FunD name1 [Clause [] (NormalB body) []] ,SigD namefull <$> [t| forall n. SNat n -> RS.Array n $ttyp -> $ttyp |] ,do body <- [| \sn -> vectorRedFullOp sn $scaleVar id id $c_opfull |] return $ FunD namefull [Clause [] (NormalB body) []] ]) $(fmap concat . forM typesList $ \arithtype -> fmap concat . forM ["min", "max"] $ \fname -> do let ttyp = conT (atType arithtype) name = mkName (fname ++ "indexVector" ++ nameBase (atType arithtype)) c_op = varE (mkName ("c_extremum_" ++ fname ++ "_" ++ atCName arithtype)) sequence [SigD name <$> [t| forall n. RS.Array n $ttyp -> [Int] |] ,do body <- [| vectorExtremumOp id $c_op |] return $ FunD name [Clause [] (NormalB body) []]]) $(fmap concat . forM typesList $ \arithtype -> do let ttyp = conT (atType arithtype) name = mkName ("dotprodinnerVector" ++ nameBase (atType arithtype)) c_op = varE (mkName ("c_dotprodinner_" ++ atCName arithtype)) mul_op = varE (mkName ("mulVector" ++ nameBase (atType arithtype))) c_scale_op = varE (mkName ("c_binary_" ++ atCName arithtype ++ "_sv")) `appE` litE (integerL (fromIntegral (aboEnum BO_MUL))) c_red_op = varE (mkName ("c_reduce1_" ++ atCName arithtype)) `appE` litE (integerL (fromIntegral (aroEnum RO_SUM))) sequence [SigD name <$> [t| forall n. SNat n -> RS.Array (n + 1) $ttyp -> RS.Array (n + 1) $ttyp -> RS.Array n $ttyp |] ,do body <- [| \sn -> vectorDotprodInnerOp sn id id $mul_op $c_scale_op $c_red_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" intWidBranchRed1 :: forall i n. (FiniteBits i, Storable i, Integral i) => -- int32 (Int64 -> Ptr Int32 -> Int32 -> Ptr Int32 -> IO ()) -- ^ scale by constant -> (Int64 -> Ptr Int32 -> Ptr Int64 -> Ptr Int64 -> 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) intWidBranchRed1 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" intWidBranchRedFull :: forall i n. (FiniteBits i, Storable i, Integral i) => (i -> Int -> i) -- ^ scale op -- int32 -> (Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int32 -> IO Int32) -- ^ reduction kernel -- int64 -> (Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int64 -> IO Int64) -- ^ reduction kernel -> (SNat n -> RS.Array n i -> i) intWidBranchRedFull fsc fred32 fred64 sn | finiteBitSize (undefined :: i) == 32 = vectorRedFullOp @i @Int32 sn fsc fromIntegral castPtr fred32 | finiteBitSize (undefined :: i) == 64 = vectorRedFullOp @i @Int64 sn fsc fromIntegral castPtr fred64 | otherwise = error "Unsupported Int width" intWidBranchExtr :: forall i n. (FiniteBits i, Storable i, Integral i) => -- int32 (Ptr Int64 -> Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int32 -> IO ()) -- ^ extremum kernel -- int64 -> (Ptr Int64 -> Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int64 -> IO ()) -- ^ extremum kernel -> (RS.Array n i -> [Int]) intWidBranchExtr fextr32 fextr64 | finiteBitSize (undefined :: i) == 32 = vectorExtremumOp @i @Int32 castPtr fextr32 | finiteBitSize (undefined :: i) == 64 = vectorExtremumOp @i @Int64 castPtr fextr64 | otherwise = error "Unsupported Int width" intWidBranchDotprod :: forall i n. (FiniteBits i, Storable i, Integral i, NumElt i) => -- int32 (Int64 -> Ptr Int32 -> Int32 -> Ptr Int32 -> IO ()) -- ^ scale by constant -> (Int64 -> Ptr Int32 -> Ptr Int64 -> Ptr Int64 -> Ptr Int32 -> IO ()) -- ^ reduction kernel -> (Int64 -> Ptr Int64 -> Ptr Int32 -> Ptr Int64 -> Ptr Int32 -> Ptr Int64 -> Ptr Int32 -> IO ()) -- ^ dotprod kernel -- int64 -> (Int64 -> Ptr Int64 -> Int64 -> Ptr Int64 -> IO ()) -- ^ scale by constant -> (Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int64 -> IO ()) -- ^ reduction kernel -> (Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int64 -> IO ()) -- ^ dotprod kernel -> (SNat n -> RS.Array (n + 1) i -> RS.Array (n + 1) i -> RS.Array n i) intWidBranchDotprod fsc32 fred32 fdot32 fsc64 fred64 fdot64 sn | finiteBitSize (undefined :: i) == 32 = vectorDotprodInnerOp @i @Int32 sn fromIntegral castPtr numEltMul fsc32 fred32 fdot32 | finiteBitSize (undefined :: i) == 64 = vectorDotprodInnerOp @i @Int64 sn fromIntegral castPtr numEltMul fsc64 fred64 fdot64 | otherwise = error "Unsupported Int width" 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 numEltSumFull :: SNat n -> RS.Array n a -> a numEltProductFull :: SNat n -> RS.Array n a -> a numEltMinIndex :: SNat n -> RS.Array n a -> [Int] numEltMaxIndex :: SNat n -> RS.Array n a -> [Int] numEltDotprodInner :: SNat n -> RS.Array (n + 1) a -> RS.Array (n + 1) a -> RS.Array n a instance NumElt Int32 where numEltAdd = addVectorInt32 numEltSub = subVectorInt32 numEltMul = mulVectorInt32 numEltNeg = negVectorInt32 numEltAbs = absVectorInt32 numEltSignum = signumVectorInt32 numEltSum1Inner = sum1VectorInt32 numEltProduct1Inner = product1VectorInt32 numEltSumFull = sumFullVectorInt32 numEltProductFull = productFullVectorInt32 numEltMinIndex _ = minindexVectorInt32 numEltMaxIndex _ = maxindexVectorInt32 numEltDotprodInner = dotprodinnerVectorInt32 instance NumElt Int64 where numEltAdd = addVectorInt64 numEltSub = subVectorInt64 numEltMul = mulVectorInt64 numEltNeg = negVectorInt64 numEltAbs = absVectorInt64 numEltSignum = signumVectorInt64 numEltSum1Inner = sum1VectorInt64 numEltProduct1Inner = product1VectorInt64 numEltSumFull = sumFullVectorInt64 numEltProductFull = productFullVectorInt64 numEltMinIndex _ = minindexVectorInt64 numEltMaxIndex _ = maxindexVectorInt64 numEltDotprodInner = dotprodinnerVectorInt64 instance NumElt Float where numEltAdd = addVectorFloat numEltSub = subVectorFloat numEltMul = mulVectorFloat numEltNeg = negVectorFloat numEltAbs = absVectorFloat numEltSignum = signumVectorFloat numEltSum1Inner = sum1VectorFloat numEltProduct1Inner = product1VectorFloat numEltSumFull = sumFullVectorFloat numEltProductFull = productFullVectorFloat numEltMinIndex _ = minindexVectorFloat numEltMaxIndex _ = maxindexVectorFloat numEltDotprodInner = dotprodinnerVectorFloat instance NumElt Double where numEltAdd = addVectorDouble numEltSub = subVectorDouble numEltMul = mulVectorDouble numEltNeg = negVectorDouble numEltAbs = absVectorDouble numEltSignum = signumVectorDouble numEltSum1Inner = sum1VectorDouble numEltProduct1Inner = product1VectorDouble numEltSumFull = sumFullVectorDouble numEltProductFull = productFullVectorDouble numEltMinIndex _ = minindexVectorDouble numEltMaxIndex _ = maxindexVectorDouble numEltDotprodInner = dotprodinnerVectorDouble instance NumElt Int where numEltAdd = intWidBranch2 @Int (+) (c_binary_i32_sv (aboEnum BO_ADD)) (flipOp (c_binary_i32_sv (aboEnum BO_ADD))) (c_binary_i32_vv (aboEnum BO_ADD)) (c_binary_i64_sv (aboEnum BO_ADD)) (flipOp (c_binary_i64_sv (aboEnum BO_ADD))) (c_binary_i64_vv (aboEnum BO_ADD)) numEltSub = intWidBranch2 @Int (-) (c_binary_i32_sv (aboEnum BO_SUB)) (flipOp (c_binary_i32_sv (aboEnum BO_SUB))) (c_binary_i32_vv (aboEnum BO_SUB)) (c_binary_i64_sv (aboEnum BO_SUB)) (flipOp (c_binary_i64_sv (aboEnum BO_SUB))) (c_binary_i64_vv (aboEnum BO_SUB)) numEltMul = intWidBranch2 @Int (*) (c_binary_i32_sv (aboEnum BO_MUL)) (flipOp (c_binary_i32_sv (aboEnum BO_MUL))) (c_binary_i32_vv (aboEnum BO_MUL)) (c_binary_i64_sv (aboEnum BO_MUL)) (flipOp (c_binary_i64_sv (aboEnum BO_MUL))) (c_binary_i64_vv (aboEnum BO_MUL)) numEltNeg = intWidBranch1 @Int (c_unary_i32 (auoEnum UO_NEG)) (c_unary_i64 (auoEnum UO_NEG)) numEltAbs = intWidBranch1 @Int (c_unary_i32 (auoEnum UO_ABS)) (c_unary_i64 (auoEnum UO_ABS)) numEltSignum = intWidBranch1 @Int (c_unary_i32 (auoEnum UO_SIGNUM)) (c_unary_i64 (auoEnum UO_SIGNUM)) numEltSum1Inner = intWidBranchRed1 @Int (c_binary_i32_sv (aboEnum BO_MUL)) (c_reduce1_i32 (aroEnum RO_SUM)) (c_binary_i64_sv (aboEnum BO_MUL)) (c_reduce1_i64 (aroEnum RO_SUM)) numEltProduct1Inner = intWidBranchRed1 @Int (c_binary_i32_sv (aboEnum BO_MUL)) (c_reduce1_i32 (aroEnum RO_PRODUCT)) (c_binary_i64_sv (aboEnum BO_MUL)) (c_reduce1_i64 (aroEnum RO_PRODUCT)) numEltSumFull = intWidBranchRedFull @Int (*) (c_reducefull_i32 (aroEnum RO_SUM)) (c_reducefull_i64 (aroEnum RO_SUM)) numEltProductFull = intWidBranchRedFull @Int (^) (c_reducefull_i32 (aroEnum RO_PRODUCT)) (c_reducefull_i64 (aroEnum RO_PRODUCT)) numEltMinIndex _ = intWidBranchExtr @Int c_extremum_min_i32 c_extremum_min_i64 numEltMaxIndex _ = intWidBranchExtr @Int c_extremum_max_i32 c_extremum_max_i64 numEltDotprodInner = intWidBranchDotprod @Int (c_binary_i32_sv (aboEnum BO_MUL)) (c_reduce1_i32 (aroEnum RO_SUM)) c_dotprodinner_i32 (c_binary_i64_sv (aboEnum BO_MUL)) (c_reduce1_i64 (aroEnum RO_SUM)) c_dotprodinner_i64 instance NumElt CInt where numEltAdd = intWidBranch2 @CInt (+) (c_binary_i32_sv (aboEnum BO_ADD)) (flipOp (c_binary_i32_sv (aboEnum BO_ADD))) (c_binary_i32_vv (aboEnum BO_ADD)) (c_binary_i64_sv (aboEnum BO_ADD)) (flipOp (c_binary_i64_sv (aboEnum BO_ADD))) (c_binary_i64_vv (aboEnum BO_ADD)) numEltSub = intWidBranch2 @CInt (-) (c_binary_i32_sv (aboEnum BO_SUB)) (flipOp (c_binary_i32_sv (aboEnum BO_SUB))) (c_binary_i32_vv (aboEnum BO_SUB)) (c_binary_i64_sv (aboEnum BO_SUB)) (flipOp (c_binary_i64_sv (aboEnum BO_SUB))) (c_binary_i64_vv (aboEnum BO_SUB)) numEltMul = intWidBranch2 @CInt (*) (c_binary_i32_sv (aboEnum BO_MUL)) (flipOp (c_binary_i32_sv (aboEnum BO_MUL))) (c_binary_i32_vv (aboEnum BO_MUL)) (c_binary_i64_sv (aboEnum BO_MUL)) (flipOp (c_binary_i64_sv (aboEnum BO_MUL))) (c_binary_i64_vv (aboEnum BO_MUL)) numEltNeg = intWidBranch1 @CInt (c_unary_i32 (auoEnum UO_NEG)) (c_unary_i64 (auoEnum UO_NEG)) numEltAbs = intWidBranch1 @CInt (c_unary_i32 (auoEnum UO_ABS)) (c_unary_i64 (auoEnum UO_ABS)) numEltSignum = intWidBranch1 @CInt (c_unary_i32 (auoEnum UO_SIGNUM)) (c_unary_i64 (auoEnum UO_SIGNUM)) numEltSum1Inner = intWidBranchRed1 @CInt (c_binary_i32_sv (aboEnum BO_MUL)) (c_reduce1_i32 (aroEnum RO_SUM)) (c_binary_i64_sv (aboEnum BO_MUL)) (c_reduce1_i64 (aroEnum RO_SUM)) numEltProduct1Inner = intWidBranchRed1 @CInt (c_binary_i32_sv (aboEnum BO_MUL)) (c_reduce1_i32 (aroEnum RO_PRODUCT)) (c_binary_i64_sv (aboEnum BO_MUL)) (c_reduce1_i64 (aroEnum RO_PRODUCT)) numEltSumFull = intWidBranchRedFull @CInt mulWithInt (c_reducefull_i32 (aroEnum RO_SUM)) (c_reducefull_i64 (aroEnum RO_SUM)) numEltProductFull = intWidBranchRedFull @CInt (^) (c_reducefull_i32 (aroEnum RO_PRODUCT)) (c_reducefull_i64 (aroEnum RO_PRODUCT)) numEltMinIndex _ = intWidBranchExtr @CInt c_extremum_min_i32 c_extremum_min_i64 numEltMaxIndex _ = intWidBranchExtr @CInt c_extremum_max_i32 c_extremum_max_i64 numEltDotprodInner = intWidBranchDotprod @CInt (c_binary_i32_sv (aboEnum BO_MUL)) (c_reduce1_i32 (aroEnum RO_SUM)) c_dotprodinner_i32 (c_binary_i64_sv (aboEnum BO_MUL)) (c_reduce1_i64 (aroEnum RO_SUM)) c_dotprodinner_i64 class FloatElt a where floatEltDiv :: SNat n -> RS.Array n a -> RS.Array n a -> RS.Array n a floatEltPow :: SNat n -> RS.Array n a -> RS.Array n a -> RS.Array n a floatEltLogbase :: SNat n -> RS.Array n a -> RS.Array n a -> RS.Array n a floatEltRecip :: SNat n -> RS.Array n a -> RS.Array n a floatEltExp :: SNat n -> RS.Array n a -> RS.Array n a floatEltLog :: SNat n -> RS.Array n a -> RS.Array n a floatEltSqrt :: SNat n -> RS.Array n a -> RS.Array n a floatEltSin :: SNat n -> RS.Array n a -> RS.Array n a floatEltCos :: SNat n -> RS.Array n a -> RS.Array n a floatEltTan :: SNat n -> RS.Array n a -> RS.Array n a floatEltAsin :: SNat n -> RS.Array n a -> RS.Array n a floatEltAcos :: SNat n -> RS.Array n a -> RS.Array n a floatEltAtan :: SNat n -> RS.Array n a -> RS.Array n a floatEltSinh :: SNat n -> RS.Array n a -> RS.Array n a floatEltCosh :: SNat n -> RS.Array n a -> RS.Array n a floatEltTanh :: SNat n -> RS.Array n a -> RS.Array n a floatEltAsinh :: SNat n -> RS.Array n a -> RS.Array n a floatEltAcosh :: SNat n -> RS.Array n a -> RS.Array n a floatEltAtanh :: SNat n -> RS.Array n a -> RS.Array n a floatEltLog1p :: SNat n -> RS.Array n a -> RS.Array n a floatEltExpm1 :: SNat n -> RS.Array n a -> RS.Array n a floatEltLog1pexp :: SNat n -> RS.Array n a -> RS.Array n a floatEltLog1mexp :: SNat n -> RS.Array n a -> RS.Array n a instance FloatElt Float where floatEltDiv = divVectorFloat floatEltPow = powVectorFloat floatEltLogbase = logbaseVectorFloat floatEltRecip = recipVectorFloat floatEltExp = expVectorFloat floatEltLog = logVectorFloat floatEltSqrt = sqrtVectorFloat floatEltSin = sinVectorFloat floatEltCos = cosVectorFloat floatEltTan = tanVectorFloat floatEltAsin = asinVectorFloat floatEltAcos = acosVectorFloat floatEltAtan = atanVectorFloat floatEltSinh = sinhVectorFloat floatEltCosh = coshVectorFloat floatEltTanh = tanhVectorFloat floatEltAsinh = asinhVectorFloat floatEltAcosh = acoshVectorFloat floatEltAtanh = atanhVectorFloat floatEltLog1p = log1pVectorFloat floatEltExpm1 = expm1VectorFloat floatEltLog1pexp = log1pexpVectorFloat floatEltLog1mexp = log1mexpVectorFloat instance FloatElt Double where floatEltDiv = divVectorDouble floatEltPow = powVectorDouble floatEltLogbase = logbaseVectorDouble floatEltRecip = recipVectorDouble floatEltExp = expVectorDouble floatEltLog = logVectorDouble floatEltSqrt = sqrtVectorDouble floatEltSin = sinVectorDouble floatEltCos = cosVectorDouble floatEltTan = tanVectorDouble floatEltAsin = asinVectorDouble floatEltAcos = acosVectorDouble floatEltAtan = atanVectorDouble floatEltSinh = sinhVectorDouble floatEltCosh = coshVectorDouble floatEltTanh = tanhVectorDouble floatEltAsinh = asinhVectorDouble floatEltAcosh = acoshVectorDouble floatEltAtanh = atanhVectorDouble floatEltLog1p = log1pVectorDouble floatEltExpm1 = expm1VectorDouble floatEltLog1pexp = log1pexpVectorDouble floatEltLog1mexp = log1mexpVectorDouble