accelerate-bignum-0.3.0.0: src/Data/Array/Accelerate/Internal/LLVM/Prim.hs
{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE TypeFamilies #-}
-- |
-- Module : Data.Array.Accelerate.Internal.LLVM.Prim
-- Copyright : [2016..2020] Trevor L. McDonell
-- License : BSD3
--
-- Maintainer : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability : experimental
-- Portability : non-portable (GHC extensions)
--
-- Primops for LLVM backends
--
module Data.Array.Accelerate.Internal.LLVM.Prim (
-- Operators from Num2
addWithCarryInt64#, mulWithCarryInt64#,
addWithCarryWord64#, mulWithCarryWord64#,
-- Operators from Num
addInt128#, subInt128#, mulInt128#,
addWord128#, subWord128#, mulWord128#,
-- Operators from Integral
quotInt128#, remInt128#, quotRemInt128#,
quotWord128#, remWord128#, quotRemWord128#,
) where
import Data.Array.Accelerate.Error
import Data.Array.Accelerate.Sugar.Elt
import Data.Array.Accelerate.Internal.BigInt
import Data.Array.Accelerate.Internal.BigWord
import Data.Array.Accelerate.Internal.Orphans.Elt ()
import Data.Array.Accelerate.LLVM.CodeGen.IR ( Operands(..) )
import Data.Array.Accelerate.LLVM.CodeGen.Monad ( CodeGen, freshName, instr_ )
import Data.Array.Accelerate.LLVM.CodeGen.Sugar
import qualified LLVM.AST.Type.Name as A
import qualified LLVM.AST.Type.Operand as A
import qualified LLVM.AST.Type.Representation as A
#if MIN_VERSION_accelerate_llvm(1,3,0)
import LLVM.AST.Type.Downcast ( downcast )
#else
import Data.Array.Accelerate.LLVM.CodeGen.Downcast ( downcast )
#endif
import LLVM.AST.Constant ( Constant(Int) )
import LLVM.AST.Instruction hiding ( nsw, nuw )
import LLVM.AST.Name ( Name(..) )
import LLVM.AST.Operand ( Operand(..) )
import LLVM.AST.Type
import Data.Int
import Data.Word
import Prelude hiding ( uncurry )
uncurry :: (Operands a -> Operands b -> c) -> Operands (((), a), b) -> c
uncurry f (OP_Unit `OP_Pair` x `OP_Pair` y) = f x y
-- Primitive instruction wrappers
-- ------------------------------
-- Operations from Num2
-- --------------------
addWithCarryInt64# :: IRFun1 arch () (EltR (Int64, Int64) -> EltR (Int64, Word64))
addWithCarryInt64# = IRFun1 $ uncurry (prim_wideningInt64 (Add nsw nuw))
mulWithCarryInt64# :: IRFun1 arch () (EltR (Int64, Int64) -> EltR (Int64, Word64))
mulWithCarryInt64# = IRFun1 $ uncurry (prim_wideningInt64 (Mul nsw nuw))
addWithCarryWord64# :: IRFun1 arch () (EltR (Word64, Word64) -> EltR (Word64, Word64))
addWithCarryWord64# = IRFun1 $ uncurry (prim_wideningWord64 (Add nsw nuw))
mulWithCarryWord64# :: IRFun1 arch () (EltR (Word64, Word64) -> EltR (Word64, Word64))
mulWithCarryWord64# = IRFun1 $ uncurry (prim_wideningWord64 (Mul nsw nuw))
prim_wideningInt64
:: (Operand -> Operand -> InstructionMetadata -> Instruction)
-> Operands Int64
-> Operands Int64
-> CodeGen arch (Operands (((), Int64), Word64))
prim_wideningInt64 op (OP_Int64 x) (OP_Int64 y) = do
a <- instr i128 (SExt (downcast x) i128 md)
b <- instr i128 (SExt (downcast y) i128 md)
c <- instr i128 (op a b md)
(d,e) <- unpackInt128 c
return $ OP_Unit `OP_Pair` upcastInt64 d `OP_Pair` upcastWord64 e
prim_wideningWord64
:: (Operand -> Operand -> InstructionMetadata -> Instruction)
-> Operands Word64
-> Operands Word64
-> CodeGen arch (Operands (((), Word64), Word64))
prim_wideningWord64 op (OP_Word64 x) (OP_Word64 y) = do
a <- instr i128 (ZExt (downcast x) i128 md)
b <- instr i128 (ZExt (downcast y) i128 md)
c <- instr i128 (op a b md)
(d,e) <- unpackWord128 c
return $ OP_Unit `OP_Pair` upcastWord64 d `OP_Pair` upcastWord64 e
-- Operations from Num
-- -------------------
addInt128# :: IRFun1 arch () (EltR (Int128, Int128) -> EltR Int128)
addInt128# = IRFun1 $ uncurry (prim_binaryInt128 (Add nsw nuw))
subInt128# :: IRFun1 arch () (EltR (Int128, Int128) -> EltR Int128)
subInt128# = IRFun1 $ uncurry (prim_binaryInt128 (Sub nsw nuw))
mulInt128# :: IRFun1 arch () (EltR (Int128, Int128) -> EltR Int128)
mulInt128# = IRFun1 $ uncurry (prim_binaryInt128 (Mul nsw nuw))
addWord128# :: IRFun1 arch () (EltR (Word128, Word128) -> EltR Word128)
addWord128# = IRFun1 $ uncurry (prim_binaryWord128 (Add nsw nuw))
subWord128# :: IRFun1 arch () (EltR (Word128, Word128) -> EltR Word128)
subWord128# = IRFun1 $ uncurry (prim_binaryWord128 (Sub nsw nuw))
mulWord128# :: IRFun1 arch () (EltR (Word128, Word128) -> EltR Word128)
mulWord128# = IRFun1 $ uncurry (prim_binaryWord128 (Mul nsw nuw))
-- Operations from Integral
-- ------------------------
quotInt128# :: IRFun1 arch () (EltR (Int128, Int128) -> EltR Int128)
quotInt128# = IRFun1 $ uncurry (prim_binaryInt128 (SDiv False))
remInt128# :: IRFun1 arch () (EltR (Int128, Int128) -> EltR Int128)
remInt128# = IRFun1 $ uncurry (prim_binaryInt128 SRem)
quotRemInt128# :: IRFun1 arch () (EltR (Int128, Int128) -> EltR (Int128, Int128))
quotRemInt128# = IRFun1 $ uncurry quotRem'
where
quotRem' :: Operands (EltR Int128) -> Operands (EltR Int128) -> CodeGen arch (Operands (EltR (Int128, Int128)))
quotRem' xx yy
| OP_Pair (OP_Pair OP_Unit (OP_Int64 xh)) (OP_Word64 xl) <- xx
, OP_Pair (OP_Pair OP_Unit (OP_Int64 yh)) (OP_Word64 yl) <- yy
= do
x <- packWord128 (downcast xh) (downcast xl)
y <- packWord128 (downcast yh) (downcast yl)
q <- instr i128 (SDiv False x y md)
z <- instr i128 (Mul nsw nuw y q md)
r <- instr i128 (Sub nsw nuw x z md)
(qh,ql) <- unpackInt128 q
(rh,rl) <- unpackInt128 r
return $ OP_Unit `OP_Pair` upcastInt128 qh ql `OP_Pair` upcastInt128 rh rl
quotWord128# :: IRFun1 arch () (EltR (Word128, Word128) -> EltR Word128)
quotWord128# = IRFun1 $ uncurry (prim_binaryWord128 (UDiv False))
remWord128# :: IRFun1 arch () (EltR (Word128, Word128) -> EltR Word128)
remWord128# = IRFun1 $ uncurry (prim_binaryWord128 URem)
quotRemWord128# :: IRFun1 arch () (EltR (Word128, Word128) -> EltR (Word128, Word128))
quotRemWord128# = IRFun1 $ uncurry quotRem'
where
quotRem' :: Operands (EltR Word128) -> Operands (EltR Word128) -> CodeGen arch (Operands (EltR (Word128, Word128)))
quotRem' xx yy
| OP_Pair (OP_Pair OP_Unit (OP_Word64 xh)) (OP_Word64 xl) <- xx
, OP_Pair (OP_Pair OP_Unit (OP_Word64 yh)) (OP_Word64 yl) <- yy
= do
x <- packWord128 (downcast xh) (downcast xl)
y <- packWord128 (downcast yh) (downcast yl)
q <- instr i128 (UDiv False x y md)
z <- instr i128 (Mul nsw nuw y q md)
r <- instr i128 (Sub nsw nuw x z md)
(qh,ql) <- unpackWord128 q
(rh,rl) <- unpackWord128 r
return $ OP_Unit `OP_Pair` upcastWord128 qh ql `OP_Pair` upcastWord128 rh rl
-- Helpers
prim_binaryWord128
:: (Operand -> Operand -> InstructionMetadata -> Instruction)
-> Operands (EltR Word128)
-> Operands (EltR Word128)
-> CodeGen arch (Operands (EltR Word128))
prim_binaryWord128 op xx yy
| OP_Pair (OP_Pair OP_Unit (OP_Word64 xh)) (OP_Word64 xl) <- xx
, OP_Pair (OP_Pair OP_Unit (OP_Word64 yh)) (OP_Word64 yl) <- yy
= do
x' <- packWord128 (downcast xh) (downcast xl)
y' <- packWord128 (downcast yh) (downcast yl)
r <- instr i128 (op x' y' md)
(hi,lo) <- unpackWord128 r
return $ upcastWord128 hi lo
prim_binaryInt128
:: (Operand -> Operand -> InstructionMetadata -> Instruction)
-> Operands (EltR Int128)
-> Operands (EltR Int128)
-> CodeGen arch (Operands (EltR Int128))
prim_binaryInt128 op xx yy
| OP_Pair (OP_Pair OP_Unit (OP_Int64 xh)) (OP_Word64 xl) <- xx
, OP_Pair (OP_Pair OP_Unit (OP_Int64 yh)) (OP_Word64 yl) <- yy
= do
x' <- packInt128 (downcast xh) (downcast xl)
y' <- packInt128 (downcast yh) (downcast yl)
r <- instr i128 (op x' y' md)
(hi,lo) <- unpackInt128 r
return $ upcastInt128 hi lo
-- Prim
-- ----
nsw :: Bool
nsw = False
nuw :: Bool
nuw = False
md :: InstructionMetadata
md = []
fresh :: CodeGen arch Name
fresh = downcast <$> freshName
instr :: Type -> Instruction -> CodeGen arch Operand
instr ty ins = do
name <- fresh
instr_ (name := ins)
return (LocalReference ty name)
packInt128 :: Operand -> Operand -> CodeGen arch Operand
packInt128 hi lo = do
a <- instr i128 (SExt hi i128 md)
b <- instr i128 (Shl nsw nuw a (ConstantOperand (Int 128 64)) md)
c <- instr i128 (ZExt lo i128 md)
d <- instr i128 (Or b c md)
return d
packWord128 :: Operand -> Operand -> CodeGen arch Operand
packWord128 hi lo = do
a <- instr i128 (ZExt hi i128 md)
b <- instr i128 (Shl nsw nuw a (ConstantOperand (Int 128 64)) md)
c <- instr i128 (ZExt lo i128 md)
d <- instr i128 (Or b c md)
return d
unpackInt128 :: Operand -> CodeGen arch (Operand,Operand)
unpackInt128 x = do
a <- instr i128 (AShr False x (ConstantOperand (Int 128 64)) md)
b <- instr i64 (Trunc a i64 md)
c <- instr i64 (Trunc x i64 md)
return (b,c)
unpackWord128 :: Operand -> CodeGen arch (Operand,Operand)
unpackWord128 x = do
a <- instr i128 (LShr False x (ConstantOperand (Int 128 64)) md)
b <- instr i64 (Trunc a i64 md)
c <- instr i64 (Trunc x i64 md)
return (b,c)
upcastInt64 :: HasCallStack => Operand -> Operands Int64
upcastInt64 (LocalReference (IntegerType 64) (UnName x)) = OP_Int64 (A.LocalReference A.type' (A.UnName x))
upcastInt64 _ = internalError "expected local reference"
upcastWord64 :: HasCallStack => Operand -> Operands Word64
upcastWord64 (LocalReference (IntegerType 64) (UnName x)) = OP_Word64 (A.LocalReference A.type' (A.UnName x))
upcastWord64 _ = internalError "expected local reference"
upcastInt128 :: Operand -> Operand -> Operands (EltR Int128)
upcastInt128 hi lo = OP_Pair (OP_Pair OP_Unit (upcastInt64 hi)) (upcastWord64 lo)
upcastWord128 :: Operand -> Operand -> Operands (EltR Word128)
upcastWord128 hi lo = OP_Pair (OP_Pair OP_Unit (upcastWord64 hi)) (upcastWord64 lo)