path: root/src/Data/Array/Mixed/Internal/Arith.hs

{-# 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, 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


-- 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 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"
  | 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 $ map fst $ filter (not . snd) (zip (zip 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) pshF pstridesR (ptrconv poutvR) (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 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 $ map fst $ filter (not . snd) (zip (zip 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

vectorDotprodOp :: (Num a, Storable a)
                => (b -> a)
                -> (Ptr a -> Ptr b)
                -> (Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr b -> Ptr b -> IO ())  -- ^ reduction kernel
                -> (Int64 -> Ptr b -> Ptr b -> IO b)  -- ^ dotprod kernel
                -> (Int64 -> Int64 -> Int64 -> Ptr b -> Int64 -> Int64 -> Ptr b -> IO b)  -- ^ strided dotprod kernel
                -> RS.Array 1 a -> RS.Array 1 a -> a
vectorDotprodOp valbackconv ptrconv fred fdot fdotstrided
    (RS.A (RG.A [len1] (OI.T [stride1] offset1 vec1)))
    (RS.A (RG.A [len2] (OI.T [stride2] offset2 vec2)))
  | len1 /= len2 = error $ "vectorDotprodOp: lengths unequal: " ++ show len1 ++ " vs " ++ show len2
  | len1 == 0 = 0  -- if the arrays are empty, just return zero
  | otherwise = case (stride1, stride2) of
      (0, 0) ->  -- replicated scalar * replicated scalar
        fromIntegral len1 * (vec1 VS.! offset1) * (vec2 VS.! offset2)
      (0, 1) ->  -- replicated scalar * dense
        dotScalarVector len1 ptrconv fred (vec1 VS.! offset1) (VS.slice offset2 len1 vec2)
      (0, -1) ->  -- replicated scalar * reversed dense
        dotScalarVector len1 ptrconv fred (vec1 VS.! offset1) (VS.slice (offset2 - (len1 - 1)) len1 vec2)
      (1, 0) ->  -- dense * replicated scalar
        dotScalarVector len1 ptrconv fred (vec2 VS.! offset2) (VS.slice offset1 len1 vec1)
      (-1, 0) ->  -- reversed dense * replicated scalar
        dotScalarVector len1 ptrconv fred (vec2 VS.! offset2) (VS.slice (offset1 - (len1 - 1)) len1 vec1)
      (1, 1) ->  -- dense * dense
        dotVectorVector len1 valbackconv ptrconv fdot (VS.slice offset1 len1 vec1) (VS.slice offset2 len1 vec2)
      (-1, -1) ->  -- reversed dense * reversed dense
        dotVectorVector len1 valbackconv ptrconv fdot (VS.slice (offset1 - (len1 - 1)) len1 vec1) (VS.slice (offset2 - (len1 - 1)) len1 vec2)
      (_, _) ->  -- fallback case
        dotVectorVectorStrided len1 valbackconv ptrconv fdotstrided offset1 stride1 vec1 offset2 stride2 vec2
vectorDotprodOp _ _ _ _ _ _ _ = error "vectorDotprodOp: not one-dimensional?"

{-# NOINLINE dotScalarVector #-}
dotScalarVector :: forall a b. (Num a, Storable a)
                => Int -> (Ptr a -> Ptr b)
                -> (Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr b -> 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 pshape pstride (ptrconv pout) (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) []]])

$(fmap concat . forM typesList $ \arithtype -> do
    let ttyp = conT (atType arithtype)
    fmap concat . forM [minBound..maxBound] $ \arithop -> do
      let name = mkName (aroName arithop ++ "Vector" ++ nameBase (atType arithtype))
          c_op = varE (mkName ("c_reduce_" ++ 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 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) []]])

$(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 ("dotprodVector" ++ nameBase (atType arithtype))
        c_op = varE (mkName ("c_dotprod_" ++ atCName arithtype))
        c_op_strided = varE (mkName ("c_dotprod_" ++ atCName arithtype ++ "_strided"))
        c_red_op = varE (mkName ("c_reduce_" ++ atCName arithtype)) `appE` litE (integerL (fromIntegral (aroEnum RO_SUM1)))
    sequence [SigD name <$>
                   [t| RS.Array 1 $ttyp -> RS.Array 1 $ttyp -> $ttyp |]
             ,do body <- [| vectorDotprodOp id id $c_red_op $c_op $c_op_strided |]
                 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"

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. (FiniteBits i, Storable i, Integral i)
                    => -- int32
                       (Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int32 -> Ptr Int32 -> IO ())  -- ^ reduction kernel
                    -> (Int64 -> Ptr Int32 -> Ptr Int32 -> IO Int32)  -- ^ dotprod kernel
                    -> (Int64 -> Int64 -> Int64 -> Ptr Int32 -> Int64 -> Int64 -> Ptr Int32 -> IO Int32)  -- ^ strided dotprod kernel
                       -- int64
                    -> (Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int64 -> Ptr Int64 -> IO ())  -- ^ reduction kernel
                    -> (Int64 -> Ptr Int64 -> Ptr Int64 -> IO Int64)  -- ^ dotprod kernel
                    -> (Int64 -> Int64 -> Int64 -> Ptr Int64 -> Int64 -> Int64 -> Ptr Int64 -> IO Int64)  -- ^ strided dotprod kernel
                    -> (RS.Array 1 i -> RS.Array 1 i -> i)
intWidBranchDotprod fred32 fdot32 fdot32strided fred64 fdot64 fdot64strided
  | finiteBitSize (undefined :: i) == 32 = vectorDotprodOp @i @Int32 fromIntegral castPtr fred32 fdot32 fdot32strided
  | finiteBitSize (undefined :: i) == 64 = vectorDotprodOp @i @Int64 fromIntegral castPtr fred64 fdot64 fdot64strided
  | 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
  numEltMinIndex :: RS.Array n a -> [Int]
  numEltMaxIndex :: RS.Array n a -> [Int]
  numEltDotprod :: RS.Array 1 a -> RS.Array 1 a -> a

instance NumElt Int32 where
  numEltAdd = addVectorInt32
  numEltSub = subVectorInt32
  numEltMul = mulVectorInt32
  numEltNeg = negVectorInt32
  numEltAbs = absVectorInt32
  numEltSignum = signumVectorInt32
  numEltSum1Inner = sum1VectorInt32
  numEltProduct1Inner = product1VectorInt32
  numEltMinIndex = minindexVectorInt32
  numEltMaxIndex = maxindexVectorInt32
  numEltDotprod = dotprodVectorInt32

instance NumElt Int64 where
  numEltAdd = addVectorInt64
  numEltSub = subVectorInt64
  numEltMul = mulVectorInt64
  numEltNeg = negVectorInt64
  numEltAbs = absVectorInt64
  numEltSignum = signumVectorInt64
  numEltSum1Inner = sum1VectorInt64
  numEltProduct1Inner = product1VectorInt64
  numEltMinIndex = minindexVectorInt64
  numEltMaxIndex = maxindexVectorInt64
  numEltDotprod = dotprodVectorInt64

instance NumElt Float where
  numEltAdd = addVectorFloat
  numEltSub = subVectorFloat
  numEltMul = mulVectorFloat
  numEltNeg = negVectorFloat
  numEltAbs = absVectorFloat
  numEltSignum = signumVectorFloat
  numEltSum1Inner = sum1VectorFloat
  numEltProduct1Inner = product1VectorFloat
  numEltMinIndex = minindexVectorFloat
  numEltMaxIndex = maxindexVectorFloat
  numEltDotprod = dotprodVectorFloat

instance NumElt Double where
  numEltAdd = addVectorDouble
  numEltSub = subVectorDouble
  numEltMul = mulVectorDouble
  numEltNeg = negVectorDouble
  numEltAbs = absVectorDouble
  numEltSignum = signumVectorDouble
  numEltSum1Inner = sum1VectorDouble
  numEltProduct1Inner = product1VectorDouble
  numEltMinIndex = minindexVectorDouble
  numEltMaxIndex = maxindexVectorDouble
  numEltDotprod = dotprodVectorDouble

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 = intWidBranchRed @Int
                      (c_binary_i32_sv (aboEnum BO_MUL)) (c_reduce_i32 (aroEnum RO_SUM1))
                      (c_binary_i64_sv (aboEnum BO_MUL)) (c_reduce_i64 (aroEnum RO_SUM1))
  numEltProduct1Inner = intWidBranchRed @Int
                          (c_binary_i32_sv (aboEnum BO_MUL)) (c_reduce_i32 (aroEnum RO_PRODUCT1))
                          (c_binary_i64_sv (aboEnum BO_MUL)) (c_reduce_i64 (aroEnum RO_PRODUCT1))
  numEltMinIndex = intWidBranchExtr @Int c_extremum_min_i32 c_extremum_min_i64
  numEltMaxIndex = intWidBranchExtr @Int c_extremum_max_i32 c_extremum_max_i64
  numEltDotprod = intWidBranchDotprod @Int (c_reduce_i32 (aroEnum RO_SUM1)) c_dotprod_i32 c_dotprod_i32_strided
                                           (c_reduce_i64 (aroEnum RO_SUM1)) c_dotprod_i64 c_dotprod_i64_strided

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 = intWidBranchRed @CInt
                      (c_binary_i32_sv (aboEnum BO_MUL)) (c_reduce_i32 (aroEnum RO_SUM1))
                      (c_binary_i64_sv (aboEnum BO_MUL)) (c_reduce_i64 (aroEnum RO_SUM1))
  numEltProduct1Inner = intWidBranchRed @CInt
                          (c_binary_i32_sv (aboEnum BO_MUL)) (c_reduce_i32 (aroEnum RO_PRODUCT1))
                          (c_binary_i64_sv (aboEnum BO_MUL)) (c_reduce_i64 (aroEnum RO_PRODUCT1))
  numEltMinIndex = intWidBranchExtr @CInt c_extremum_min_i32 c_extremum_min_i64
  numEltMaxIndex = intWidBranchExtr @CInt c_extremum_max_i32 c_extremum_max_i64
  numEltDotprod = intWidBranchDotprod @CInt (c_reduce_i32 (aroEnum RO_SUM1)) c_dotprod_i32 c_dotprod_i32_strided
                                            (c_reduce_i64 (aroEnum RO_SUM1)) c_dotprod_i64 c_dotprod_i64_strided

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