accelerate-cuda-0.17.0.0: Data/Array/Accelerate/CUDA/CodeGen/Arithmetic.hs
{-# LANGUAGE GADTs #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ViewPatterns #-}
-- |
-- Module : Data.Array.Accelerate.CUDA.CodeGen.Arithmetic
-- Copyright : [2008..2014] Manuel M T Chakravarty, Gabriele Keller
-- [2009..2014] Trevor L. McDonell
-- License : BSD3
--
-- Maintainer : Trevor L. McDonell <tmcdonell@cse.unsw.edu.au>
-- Stability : experimental
-- Portability : non-portable (GHC extensions)
--
module Data.Array.Accelerate.CUDA.CodeGen.Arithmetic
where
-- friends
import Data.Array.Accelerate.Type
import Data.Array.Accelerate.Error
import Data.Array.Accelerate.CUDA.CodeGen.Base
import Data.Array.Accelerate.CUDA.CodeGen.Constant
import Data.Array.Accelerate.CUDA.CodeGen.Monad
import Data.Array.Accelerate.CUDA.CodeGen.Type
-- libraries
import Prelude ( String, ($), (++), (-), undefined, otherwise )
import Data.Bits ( finiteBitSize )
import Control.Monad.State.Strict
import Language.C
import Language.C.Quote.CUDA
import Foreign.Storable ( sizeOf )
-- Operations from Num
-- ===================
add :: Exp -> Exp -> Exp
add x y = [cexp| $exp:x + $exp:y |]
sub :: Exp -> Exp -> Exp
sub x y = [cexp| $exp:x - $exp:y |]
mul :: Exp -> Exp -> Exp
mul x y = [cexp| $exp:x * $exp:y |]
negate :: Exp -> Exp
negate x = [cexp| - $exp:x |]
abs :: forall a. NumType a -> Exp -> Exp
abs (FloatingNumType t) x
= mathf t "fabs" [x]
abs (IntegralNumType t) x
| signedIntegralNum t
, IntegralDict <- integralDict t
= case sizeOf (undefined::a) of
8 -> ccall "llabs" [x]
_ -> ccall "abs" [x]
| otherwise
= x
signum :: NumType a -> Exp -> Gen Exp
signum (IntegralNumType t) x
| IntegralDict <- integralDict t
, unsignedIntegralNum t
= return [cexp| $exp:x > $exp:(integral t 0) |]
| IntegralDict <- integralDict t
= do x' <- bind (typeOf t) x
return [cexp| ($exp:x' > $exp:(integral t 0)) - ($exp:x' < $exp:(integral t 0)) |]
signum (FloatingNumType t) x
| FloatingDict <- floatingDict t
= do x' <- bind (typeOf t) x
return [cexp| $exp:x' == $exp:(floating t 0)
? $exp:(floating t 0)
: $exp:(mathf t "copysign" [floating t 1, x']) |]
-- Operators from Integral & Bits
-- ==============================
quot :: Exp -> Exp -> Exp
quot x y = [cexp| $exp:x / $exp:y |]
rem :: Exp -> Exp -> Exp
rem x y = [cexp| $exp:x % $exp:y |]
quotRem :: IntegralType a -> Exp -> Exp -> Gen (Exp,Exp)
quotRem (typeOf -> t') x y = do
x' <- bind t' x
y' <- bind t' y
q <- bind t' (x' `quot` y')
r <- bind t' (x' `sub` (y' `mul` q))
return (q, r)
idiv :: IntegralType a -> Exp -> Exp -> Gen Exp
idiv t x y
| unsignedIntegralNum t
= return (x `quot` y)
| IntegralDict <- integralDict t
, zero <- integral t 0
, one <- integral t 1
= do
x' <- bind (typeOf t) x
y' <- bind (typeOf t) y
return $
cases [ ((x' `gt` zero) `land` (y' `lt` zero), ((x' `sub` one) `quot` y') `sub` one)
, ((x' `lt` zero) `land` (y' `gt` zero), ((x' `add` one) `quot` y') `sub` one)
]
(x' `quot` y')
mod :: IntegralType a -> Exp -> Exp -> Gen Exp
mod t x y
| unsignedIntegralNum t
= return (x `rem` y)
| IntegralDict <- integralDict t
, zero <- integral t 0
= do
x' <- bind (typeOf t) x
y' <- bind (typeOf t) y
r <- bind (typeOf t) (x' `rem` y')
return $
((((x' `gt` zero) `land` (y' `lt` zero)) `lor` ((x' `lt` zero) `land` (y' `gt` zero)))
?: ( r `neq` zero ?: ( r `add` y', zero )
, r ))
divMod :: IntegralType a -> Exp -> Exp -> Gen (Exp, Exp)
divMod t x y | IntegralDict <- integralDict t = do
x' <- bind (typeOf t) x
y' <- bind (typeOf t) y
(q,r) <- quotRem t x' y'
sr <- signum (IntegralNumType t) r
sy' <- signum (IntegralNumType t) y'
-- Somewhat awful way to inject an ifThenElse statement
vd <- lift fresh
vm <- lift fresh
modify (\st -> st { localBindings = [citem| $ty:(typeOf t) $id:vd, $id:vm; |] : localBindings st })
modify (\st -> st { localBindings = [citem| if ( $exp:(sr `eq` negate sy') ) {
$id:vd = $exp:(q `sub` integral t 1);
$id:vm = $exp:(r `add` y') ;
} else {
$id:vd = $exp:q;
$id:vm = $exp:r;
} |] : localBindings st })
return ( cvar vd, cvar vm )
band :: Exp -> Exp -> Exp
band x y = [cexp| $exp:x & $exp:y |]
bor :: Exp -> Exp -> Exp
bor x y = [cexp| $exp:x | $exp:y |]
xor :: Exp -> Exp -> Exp
xor x y = [cexp| $exp:x ^ $exp:y |]
bnot :: Exp -> Exp
bnot x = [cexp| ~ $exp:x |]
shiftL :: Exp -> Exp -> Exp
shiftL x i = [cexp| $exp:x << $exp:i |]
-- Arithmetic right shift (unchecked)
--
shiftRA :: Exp -> Exp -> Exp
shiftRA x i = [cexp| $exp:x >> $exp:i |]
-- Logical right shift (unchecked)
--
shiftRL :: IntegralType a -> Exp -> Exp -> Exp
shiftRL ty x i =
let int = typeOf (integralType :: IntegralType Int)
word = typeOf (integralType :: IntegralType Word)
in
case ty of
TypeInt{} -> [cexp| ($ty:int) (($ty:word) $exp:x >> $exp:i) |]
TypeInt8{} -> [cexp| (typename Int8) ((typename Word8) $exp:x >> $exp:i) |]
TypeInt16{} -> [cexp| (typename Int16) ((typename Word16) $exp:x >> $exp:i) |]
TypeInt32{} -> [cexp| (typename Int32) ((typename Word32) $exp:x >> $exp:i) |]
TypeInt64{} -> [cexp| (typename Int64) ((typename Word64) $exp:x >> $exp:i) |]
TypeCShort{} -> [cexp| (short) ((unsigned short) $exp:x >> $exp:i) |]
TypeCInt{} -> [cexp| (int) ((unsigned int) $exp:x >> $exp:i) |]
TypeCLong{} -> [cexp| (long) ((unsigned long) $exp:x >> $exp:i) |]
TypeCLLong{} -> [cexp| (long long) ((unsigned long long) $exp:x >> $exp:i) |]
-- unsigned types use arithmetic shift
_ -> $internalCheck "shiftRL" "unhandled signed type" (unsignedIntegralNum ty) (shiftRA x i)
rotateL :: forall a. IntegralType a -> Exp -> Exp -> Gen Exp
rotateL t x i | IntegralDict <- integralDict t = do
let int = integralType :: IntegralType Int
wsib = finiteBitSize (undefined::a)
--
x' <- bind (typeOf t) x
i' <- bind (typeOf int) (i `band` integral int (wsib - 1))
return $ (x' `shiftL` i') `bor` (shiftRL t x' (integral int wsib `sub` i'))
rotateR :: IntegralType a -> Exp -> Exp -> Gen Exp
rotateR t x i = rotateL t x (negate i)
-- Operators from Fractional & Floating
-- ====================================
fdiv :: Exp -> Exp -> Exp
fdiv x y = [cexp| $exp:x / $exp:y |]
recip :: FloatingType a -> Exp -> Exp
recip t x | FloatingDict <- floatingDict t = fdiv (floating t 1) x
sin :: FloatingType a -> Exp -> Exp
sin t x = mathf t "sin" [x]
cos :: FloatingType a -> Exp -> Exp
cos t x = mathf t "cos" [x]
tan :: FloatingType a -> Exp -> Exp
tan t x = mathf t "tan" [x]
asin :: FloatingType a -> Exp -> Exp
asin t x = mathf t "asin" [x]
acos :: FloatingType a -> Exp -> Exp
acos t x = mathf t "acos" [x]
atan :: FloatingType a -> Exp -> Exp
atan t x = mathf t "atan" [x]
sinh :: FloatingType a -> Exp -> Exp
sinh t x = mathf t "sinh" [x]
cosh :: FloatingType a -> Exp -> Exp
cosh t x = mathf t "cosh" [x]
tanh :: FloatingType a -> Exp -> Exp
tanh t x = mathf t "tanh" [x]
asinh :: FloatingType a -> Exp -> Exp
asinh t x = mathf t "asinh" [x]
acosh :: FloatingType a -> Exp -> Exp
acosh t x = mathf t "acosh" [x]
atanh :: FloatingType a -> Exp -> Exp
atanh t x = mathf t "atanh" [x]
exp :: FloatingType a -> Exp -> Exp
exp t x = mathf t "exp" [x]
sqrt :: FloatingType a -> Exp -> Exp
sqrt t x = mathf t "sqrt" [x]
pow :: FloatingType a -> Exp -> Exp -> Exp
pow t x y = mathf t "pow" [x,y]
log :: FloatingType a -> Exp -> Exp
log t x = mathf t "log" [x]
logBase :: FloatingType a -> Exp -> Exp -> Exp
logBase t x y = log t y `fdiv` log t x
-- Operators from RealFrac
-- =======================
trunc :: FloatingType a -> IntegralType b -> Exp -> Exp
trunc ta tb x = cast tb $ mathf ta "trunc" [x]
round :: FloatingType a -> IntegralType b -> Exp -> Exp
round ta tb x = cast tb $ mathf ta "round" [x]
floor :: FloatingType a -> IntegralType b -> Exp -> Exp
floor ta tb x = cast tb $ mathf ta "floor" [x]
ceiling :: FloatingType a -> IntegralType b -> Exp -> Exp
ceiling ta tb x = cast tb $ mathf ta "ceil" [x]
-- Operators from RealFloat
-- ========================
atan2 :: FloatingType a -> Exp -> Exp -> Exp
atan2 t x y = mathf t "atan2" [x, y]
isNaN :: Exp -> Exp
isNaN x = ccall "isnan" [x]
-- Relational and equality operators
-- =================================
lt :: Exp -> Exp -> Exp
lt x y = [cexp| $exp:x < $exp:y |]
gt :: Exp -> Exp -> Exp
gt x y = [cexp| $exp:x > $exp:y |]
leq :: Exp -> Exp -> Exp
leq x y = [cexp| $exp:x <= $exp:y |]
geq :: Exp -> Exp -> Exp
geq x y = [cexp| $exp:x >= $exp:y |]
eq :: Exp -> Exp -> Exp
eq x y = [cexp| $exp:x == $exp:y |]
neq :: Exp -> Exp -> Exp
neq x y = [cexp| $exp:x != $exp:y |]
max :: ScalarType a -> Exp -> Exp -> Exp
max (NonNumScalarType _) x y =
let t = scalarType :: ScalarType Int32
in max t (cast t x) (cast t y)
max (NumScalarType (IntegralNumType _)) x y = ccall "max" [x,y]
max (NumScalarType (FloatingNumType t)) x y = mathf t "fmax" [x,y]
min :: ScalarType a -> Exp -> Exp -> Exp
min (NonNumScalarType _) x y =
let t = scalarType :: ScalarType Int32
in min t (cast t x) (cast t y)
min (NumScalarType (IntegralNumType _)) x y = ccall "min" [x,y]
min (NumScalarType (FloatingNumType t)) x y = mathf t "fmin" [x,y]
-- Logical operators
-- =================
land :: Exp -> Exp -> Exp
land x y = [cexp| $exp:x && $exp:y |]
lor :: Exp -> Exp -> Exp
lor x y = [cexp| $exp:x || $exp:y |]
lnot :: Exp -> Exp
lnot x = [cexp| ! $exp:x |]
-- Type Conversions
-- ================
ord :: Exp -> Exp
ord = cast (scalarType :: ScalarType Int)
chr :: Exp -> Exp
chr = cast (scalarType :: ScalarType Char)
boolToInt :: Exp -> Exp
boolToInt = cast (scalarType :: ScalarType Int)
fromIntegral :: IntegralType a -> NumType b -> Exp -> Exp
fromIntegral _ tb = cast tb
-- Helpers
-- =======
cast :: TypeOf a => a -> Exp -> Exp
cast t x = [cexp| ($ty:(typeOf t)) $exp:x |]
mathf :: forall t. FloatingType t -> String -> [Exp] -> Exp
mathf ty f args | FloatingDict <- floatingDict ty =
let
fun = f ++ case sizeOf (undefined :: t) of
4 -> "f"
8 -> []
16 -> "l" -- long double
_ -> $internalError "mathf" "unsupported floating point size"
in
ccall fun args
infix 0 ?:
(?:) :: Exp -> (Exp, Exp) -> Exp
(?:) p (t,e) = [cexp| $exp:p ? $exp:t : $exp:e |]
cases :: [(Exp, Exp)] -> Exp -> Exp
cases [] def = def
cases ((p,b):rest) def = p ?: (b, cases rest def)