llvm-dsl-0.2.1: src/LLVM/DSL/Expression/Vector.hs

{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module LLVM.DSL.Expression.Vector where

import qualified LLVM.DSL.Expression as Expr
import LLVM.DSL.Expression (Exp)

import qualified LLVM.Extra.Nice.Value.Vector as NiceValueVec
import qualified LLVM.Extra.Nice.Value as NiceValue
import qualified LLVM.Extra.Nice.Vector.Instance as NiceVectorInst
import qualified LLVM.Extra.Nice.Vector as NiceVector
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Core as LLVM

import qualified Data.Tuple.HT as Tuple

import Prelude hiding (replicate, take, zip, fst, snd, min, max)


{- $setup
>>> import Test.LLVM.DSL.Run (sharedRun, run)
>>> import qualified LLVM.DSL.Expression.Vector as Vector
>>> import qualified LLVM.DSL.Expression as Expr
>>> import qualified LLVM.Core as LLVM
>>> import Type.Data.Num.Decimal (D4, D8)
>>> import qualified Data.Traversable as Trav
>>> import qualified Data.Foldable as Fold
>>> import qualified Data.List as List
>>> import qualified Data.NonEmpty as NonEmpty
>>> import qualified Data.Tuple.HT as Tuple
>>> import Control.Applicative (liftA2)
>>> import Data.Int (Int32)
>>> import Test.QuickCheck ((===))
>>>
>>> type N = D8
>>> type VectorN = LLVM.Vector N
-}


cons ::
   (LLVM.Positive n, NiceVector.C a) =>
   LLVM.Vector n a -> Exp (LLVM.Vector n a)
cons = Expr.lift0 . NiceValueVec.cons

fst ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b) =>
   Exp (LLVM.Vector n (a,b)) -> Exp (LLVM.Vector n a)
fst = Expr.lift1 NiceValueVec.fst

snd ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b) =>
   Exp (LLVM.Vector n (a,b)) -> Exp (LLVM.Vector n b)
snd = Expr.lift1 NiceValueVec.snd

swap ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b) =>
   Exp (LLVM.Vector n (a,b)) -> Exp (LLVM.Vector n (b,a))
swap = Expr.lift1 NiceValueVec.swap

mapFst ::
   (Exp (LLVM.Vector n a0) -> Exp (LLVM.Vector n a1)) ->
   Exp (LLVM.Vector n (a0,b)) -> Exp (LLVM.Vector n (a1,b))
mapFst f =
   Expr.liftReprM
      (\(a0,b) -> do
         NiceValue.Cons a1 <- Expr.unliftM1 f $ NiceValue.Cons a0
         return (a1,b))

mapSnd ::
   (Exp (LLVM.Vector n b0) -> Exp (LLVM.Vector n b1)) ->
   Exp (LLVM.Vector n (a,b0)) -> Exp (LLVM.Vector n (a,b1))
mapSnd f =
   Expr.liftReprM
      (\(a,b0) -> do
         NiceValue.Cons b1 <- Expr.unliftM1 f $ NiceValue.Cons b0
         return (a,b1))


fst3 ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b, NiceVector.C c) =>
   Exp (LLVM.Vector n (a,b,c)) -> Exp (LLVM.Vector n a)
fst3 = Expr.lift1 NiceValueVec.fst3

snd3 ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b, NiceVector.C c) =>
   Exp (LLVM.Vector n (a,b,c)) -> Exp (LLVM.Vector n b)
snd3 = Expr.lift1 NiceValueVec.snd3

thd3 ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b, NiceVector.C c) =>
   Exp (LLVM.Vector n (a,b,c)) -> Exp (LLVM.Vector n c)
thd3 = Expr.lift1 NiceValueVec.thd3


zip ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b) =>
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n b) ->
   Exp (LLVM.Vector n (a,b))
zip = Expr.lift2 NiceValueVec.zip

zip3 ::
   (LLVM.Positive n, NiceVector.C a, NiceVector.C b, NiceVector.C c) =>
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n b) -> Exp (LLVM.Vector n c) ->
   Exp (LLVM.Vector n (a,b,c))
zip3 = Expr.lift3 NiceValueVec.zip3


{- |
>>> run Vector.replicate 42 :: LLVM.Vector D4 Int32
Vector [42,42,42,42]
-}
replicate ::
   (LLVM.Positive n, NiceVector.C a) =>
   Exp a -> Exp (LLVM.Vector n a)
replicate = Expr.liftM NiceValueVec.replicate

{- |
>>> run (Expr.uncurry3 Vector.replicateSegmented) (LLVM.consVector False False True False False True False False, 8, LLVM.consVector 7 6 5 4 3 2 1 0) :: LLVM.Vector D8 Int32
Vector [8,8,5,5,5,2,2,2]

prop> :{
   sharedRun (Expr.uncurry3 Vector.replicateSegmented) $ \replicateSegmented ->
   \bs x xs ->
      NonEmpty.tail
         (NonEmpty.scanl (\lastX (b,xi) -> if b then xi else lastX)
            x (liftA2 (,) bs xs))
      ===
      replicateSegmented (bs, x, xs :: VectorN Int32)
:}
-}
replicateSegmented ::
   (LLVM.Positive n, NiceVector.Select a) =>
   Exp (LLVM.Vector n Bool) -> Exp a -> Exp (LLVM.Vector n a) ->
   Exp (LLVM.Vector n a)
replicateSegmented = Expr.liftM3 NiceValueVec.replicateSegmented

{- |
>>> run (Vector.iterate (6+)) 0 :: LLVM.Vector D8 Int32
Vector [0,6,12,18,24,30,36,42]

prop> :{
   sharedRun (Expr.uncurry (\d -> Vector.iterate (d+))) $ \iter ->
   \d x ->
      List.take 8 (List.iterate (d+) x)
      ===
      Fold.toList (iter (d,x) :: LLVM.Vector D8 Int32)
:}
-}
iterate ::
   (LLVM.Positive n, NiceVector.C a) =>
   (Exp a -> Exp a) -> Exp a -> Exp (LLVM.Vector n a)
iterate f = Expr.liftM (NiceValueVec.iterate (Expr.unliftM1 f))

take ::
   (LLVM.Positive n, LLVM.Positive m, NiceVector.Select a) =>
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector m a)
take = Expr.liftM NiceValueVec.take

takeRev ::
   (LLVM.Positive n, LLVM.Positive m, NiceVector.Select a) =>
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector m a)
takeRev = Expr.liftM NiceValueVec.takeRev


{- |
Same as 'cumulateJoint' but misses sharing.
That is, cumulation operations is effictively computed twice,
unless LLVM optimization eliminates it.
However, interim code can blow up
if many such unshared operations are nested.
-}
cumulate ::
   (LLVM.Positive n, NiceVector.Additive a) =>
   Exp a -> Exp (LLVM.Vector n a) -> (Exp a, Exp (LLVM.Vector n a))
cumulate a v = Expr.unzip $ cumulateJoint a v

{- |
>>> run (Expr.uncurry Vector.cumulateJoint) (3, LLVM.consVector 1 4 1 5 :: LLVM.Vector D4 Int32)
(14,Vector [3,4,8,9])

prop> :{
   sharedRun (Expr.uncurry Vector.cumulateJoint) $ \cumulateJoint ->
   \x xs ->
      Tuple.swap (NonEmpty.viewR (NonEmpty.scanl (+) x xs))
      ===
      cumulateJoint (x, xs :: VectorN Int32)
:}
-}
cumulateJoint ::
   (LLVM.Positive n, NiceVector.Additive a) =>
   Exp a -> Exp (LLVM.Vector n a) -> Exp (a, LLVM.Vector n a)
cumulateJoint a0 v0 =
   Expr.liftM2
      (\a v ->
         fmap (uncurry NiceValue.zip .
               Tuple.mapSnd NiceVectorInst.toNiceValue) $
         NiceVector.cumulate a $ NiceVectorInst.fromNiceValue v)
      a0 v0

{- |
>>> run Vector.cumulate1 (LLVM.consVector 3 1 4 1 :: LLVM.Vector D4 Int32)
Vector [3,4,8,9]

>>> run Vector.cumulate1 (LLVM.consVector 3 1 4 1 5 9 2 6 :: LLVM.Vector D8 Int32)
Vector [3,4,8,9,14,23,25,31]

>>> run Vector.cumulate1 (LLVM.consVector 0 0 1 0 0 (-1) 0 0 :: LLVM.Vector D8 Int32)
Vector [0,0,1,1,1,0,0,0]

prop> :{
   sharedRun Vector.cumulate1 $ \cumulate1 ->
   \x ->
      scanl1 (+) (Fold.toList x)
      ===
      Fold.toList (cumulate1 (x :: VectorN Int32))
:}

prop> :{
   sharedRun Vector.cumulate1 $ \cumulate1 ->
   \x ->
      NonEmpty.tail (NonEmpty.scanl (+) 0 x)
      ===
      cumulate1 (x :: VectorN Int32)
:}
-}
cumulate1 ::
   (LLVM.Positive n, NiceVector.Additive a) =>
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n a)
cumulate1 =
   Expr.liftM
      (fmap NiceVectorInst.toNiceValue .
       NiceVector.cumulate1 .
       NiceVectorInst.fromNiceValue)

cumulateSegmented ::
   (LLVM.Positive n, NiceVector.Additive a, NiceVector.Select a) =>
   Exp (LLVM.Vector n Bool) -> Exp a -> Exp (LLVM.Vector n a) ->
   (Exp a, Exp (LLVM.Vector n a))
cumulateSegmented b a v = Expr.unzip $ cumulateSegmentedJoint b a v

{- |
>>> run (Expr.uncurry3 Vector.cumulateSegmentedJoint) (LLVM.consVector False False True False False True False False, 3, LLVM.consVector 1 4 1 5 9 2 6 5) :: (Int32, LLVM.Vector D8 Int32)
(13,Vector [3,4,0,1,6,0,2,8])

prop> :{
   sharedRun (Expr.uncurry3 Vector.cumulateSegmentedJoint) $
   \cumulateSegmentedJoint ->
   \bs x xs ->
      Trav.mapAccumL
         (\acc (b,xi) ->
            let yi = if b then 0 else acc in (yi+xi, yi))
         x (liftA2 (,) bs xs)
      ===
      cumulateSegmentedJoint (bs, x, xs :: VectorN Int32)
:}
-}
cumulateSegmentedJoint ::
   (LLVM.Positive n, NiceVector.Additive a, NiceVector.Select a) =>
   Exp (LLVM.Vector n Bool) -> Exp a -> Exp (LLVM.Vector n a) ->
   Exp (a, LLVM.Vector n a)
cumulateSegmentedJoint b0 a0 v0 =
   Expr.liftM3
      (\b a v ->
         fmap (uncurry NiceValue.zip .
               Tuple.mapSnd NiceVectorInst.toNiceValue) $
         NiceVector.cumulateSegmented
            (NiceVectorInst.fromNiceValue b)
            a
            (NiceVectorInst.fromNiceValue v))
      b0 a0 v0


cmp ::
   (LLVM.Positive n, NiceVector.Comparison a) =>
   LLVM.CmpPredicate ->
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n Bool)
cmp ord = Expr.liftM2 (NiceValueVec.cmp ord)

select ::
   (LLVM.Positive n, NiceVector.Select a) =>
   Exp (LLVM.Vector n Bool) ->
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n a)
select = Expr.liftM3 NiceValueVec.select


min, max ::
   (LLVM.Positive n, NiceVector.Real a) =>
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n a)
min = Expr.liftM2 A.min
max = Expr.liftM2 A.max

limit ::
   (LLVM.Positive n, NiceVector.Real a) =>
   (Exp (LLVM.Vector n a), Exp (LLVM.Vector n a)) ->
   Exp (LLVM.Vector n a) -> Exp (LLVM.Vector n a)
limit (l,u) = max l . min u


fromIntegral ::
   (NiceValueVec.NativeInteger i ir, NiceValueVec.NativeFloating a ar,
    LLVM.ShapeOf ir ~ LLVM.ShapeOf ar) =>
   Exp i -> Exp a
fromIntegral = Expr.liftM NiceValueVec.fromIntegral

truncateToInt ::
   (NiceValueVec.NativeInteger i ir, NiceValueVec.NativeFloating a ar,
    LLVM.ShapeOf ir ~ LLVM.ShapeOf ar) =>
   Exp a -> Exp i
truncateToInt = Expr.liftM NiceValueVec.truncateToInt

splitFractionToInt ::
   (NiceValueVec.NativeInteger i ir, NiceValueVec.NativeFloating a ar,
    LLVM.ShapeOf ir ~ LLVM.ShapeOf ar) =>
   Exp a -> (Exp i, Exp a)
splitFractionToInt = Expr.unzip . Expr.liftM NiceValueVec.splitFractionToInt