accelerate-llvm-1.4.0.0: src/Data/Array/Accelerate/LLVM/CodeGen/Exp.hs
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_HADDOCK hide #-}
-- |
-- Module : Data.Array.Accelerate.LLVM.CodeGen.Exp
-- Copyright : [2015..2020] The Accelerate Team
-- License : BSD3
--
-- Maintainer : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability : experimental
-- Portability : non-portable (GHC extensions)
--
module Data.Array.Accelerate.LLVM.CodeGen.Exp
where
import Data.Array.Accelerate.AST
import Data.Array.Accelerate.AST.LeftHandSide
import Data.Array.Accelerate.AST.Var
import Data.Array.Accelerate.Analysis.Match
import Data.Array.Accelerate.Error
import Data.Array.Accelerate.Representation.Array ( Array, arrayRshape )
import Data.Array.Accelerate.Representation.Shape
import Data.Array.Accelerate.Representation.Slice
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Representation.Vec
import Data.Array.Accelerate.Type
import qualified Data.Array.Accelerate.Sugar.Foreign as A
import Data.Array.Accelerate.LLVM.CodeGen.Array
import Data.Array.Accelerate.LLVM.CodeGen.Base
import Data.Array.Accelerate.LLVM.CodeGen.Constant
import Data.Array.Accelerate.LLVM.CodeGen.Environment
import Data.Array.Accelerate.LLVM.CodeGen.IR
import Data.Array.Accelerate.LLVM.CodeGen.Monad
import Data.Array.Accelerate.LLVM.CodeGen.Sugar
import Data.Array.Accelerate.LLVM.Foreign
import qualified Data.Array.Accelerate.LLVM.CodeGen.Arithmetic as A
import qualified Data.Array.Accelerate.LLVM.CodeGen.Loop as L
import Data.Primitive.Vec
import LLVM.AST.Type.Instruction
import LLVM.AST.Type.Operand ( Operand )
import Control.Applicative hiding ( Const )
import Control.Monad
import Prelude hiding ( exp, any )
import qualified Data.IntMap as IM
import GHC.TypeNats
-- Scalar expressions
-- ==================
{-# INLINEABLE llvmOfFun1 #-}
llvmOfFun1
:: (HasCallStack, Foreign arch)
=> Fun aenv (a -> b)
-> Gamma aenv
-> IRFun1 arch aenv (a -> b)
llvmOfFun1 (Lam lhs (Body body)) aenv = IRFun1 $ \x -> llvmOfOpenExp body (Empty `pushE` (lhs, x)) aenv
llvmOfFun1 _ _ = internalError "impossible evaluation"
{-# INLINEABLE llvmOfFun2 #-}
llvmOfFun2
:: (HasCallStack, Foreign arch)
=> Fun aenv (a -> b -> c)
-> Gamma aenv
-> IRFun2 arch aenv (a -> b -> c)
llvmOfFun2 (Lam lhs1 (Lam lhs2 (Body body))) aenv = IRFun2 $ \x y -> llvmOfOpenExp body (Empty `pushE` (lhs1, x) `pushE` (lhs2, y)) aenv
llvmOfFun2 _ _ = internalError "impossible evaluation"
-- | Convert an open scalar expression into a sequence of LLVM Operands instructions.
-- Code is generated in depth first order, and uses a monad to collect the
-- sequence of instructions used to construct basic blocks.
--
{-# INLINEABLE llvmOfOpenExp #-}
llvmOfOpenExp
:: forall arch env aenv _t. (HasCallStack, Foreign arch)
=> OpenExp env aenv _t
-> Val env
-> Gamma aenv
-> IROpenExp arch env aenv _t
llvmOfOpenExp top env aenv = cvtE top
where
cvtF1 :: OpenFun env aenv (a -> b) -> IROpenFun1 arch env aenv (a -> b)
cvtF1 (Lam lhs (Body body)) = IRFun1 $ \x -> llvmOfOpenExp body (env `pushE` (lhs, x)) aenv
cvtF1 _ = internalError "impossible evaluation"
cvtE :: forall t. OpenExp env aenv t -> IROpenExp arch env aenv t
cvtE exp =
case exp of
Let lhs bnd body -> do x <- cvtE bnd
llvmOfOpenExp body (env `pushE` (lhs, x)) aenv
Evar (Var _ ix) -> return $ prj ix env
Const tp c -> return $ ir tp $ scalar tp c
PrimConst c -> let tp = SingleScalarType (primConstType c)
in return $ ir tp $ scalar tp $ primConst c
PrimApp f x -> primFun f x
Undef tp -> return $ ir tp $ undef tp
Nil -> return $ OP_Unit
Pair e1 e2 -> join $ pair <$> cvtE e1 <*> cvtE e2
VecPack vecr e -> vecPack vecr =<< cvtE e
VecUnpack vecr e -> vecUnpack vecr =<< cvtE e
Foreign tp asm f x -> foreignE tp asm f =<< cvtE x
Case _ [] _ -> internalError "Empty Case"
Case tag xs@((_, e1):_) mx -> A.caseof (expType e1) (cvtE tag) [(t,cvtE e) | (t,e) <- xs] (fmap cvtE mx)
Cond c t e -> cond (expType t) (cvtE c) (cvtE t) (cvtE e)
IndexSlice slice slix sh -> indexSlice slice <$> cvtE slix <*> cvtE sh
IndexFull slice slix sh -> indexFull slice <$> cvtE slix <*> cvtE sh
ToIndex shr sh ix -> join $ intOfIndex shr <$> cvtE sh <*> cvtE ix
FromIndex shr sh ix -> join $ indexOfInt shr <$> cvtE sh <*> cvtE ix
Index acc ix -> index acc =<< cvtE ix
LinearIndex acc ix -> linearIndex acc =<< cvtE ix
ShapeSize shr sh -> shapeSize shr =<< cvtE sh
Shape acc -> return $ shape acc
While c f x -> while (expType x) (cvtF1 c) (cvtF1 f) (cvtE x)
Coerce t1 t2 x -> coerce t1 t2 =<< cvtE x
indexSlice :: SliceIndex slix sl co sh -> Operands slix -> Operands sh -> Operands sl
indexSlice SliceNil OP_Unit OP_Unit = OP_Unit
indexSlice (SliceAll sliceIdx) (OP_Pair slx OP_Unit) (OP_Pair sl sz) =
let sl' = indexSlice sliceIdx slx sl
in OP_Pair sl' sz
indexSlice (SliceFixed sliceIdx) (OP_Pair slx _i) (OP_Pair sl _sz) =
indexSlice sliceIdx slx sl
indexFull :: SliceIndex slix sl co sh -> Operands slix -> Operands sl -> Operands sh
indexFull SliceNil OP_Unit OP_Unit = OP_Unit
indexFull (SliceAll sliceIdx) (OP_Pair slx OP_Unit) (OP_Pair sl sz) =
let sh' = indexFull sliceIdx slx sl
in OP_Pair sh' sz
indexFull (SliceFixed sliceIdx) (OP_Pair slx sz) sl =
let sh' = indexFull sliceIdx slx sl
in OP_Pair sh' sz
vecPack :: forall n single tuple. (HasCallStack, KnownNat n) => VecR n single tuple -> Operands tuple -> CodeGen arch (Operands (Vec n single))
vecPack vecr tuple = ir tp <$> go vecr n tuple
where
go :: VecR n' single tuple' -> Int -> Operands tuple' -> CodeGen arch (Operand (Vec n single))
go (VecRnil _) 0 OP_Unit = return $ undef $ VectorScalarType tp
go (VecRnil _) _ OP_Unit = internalError "index mismatch"
go (VecRsucc vecr') i (OP_Pair xs x) = do
vec <- go vecr' (i - 1) xs
instr' $ InsertElement (fromIntegral i - 1) vec (op singleTp x)
singleTp :: SingleType single -- GHC 8.4 cannot infer this type for some reason
tp@(VectorType n singleTp) = vecRvector vecr
vecUnpack :: forall n single tuple. (HasCallStack, KnownNat n) => VecR n single tuple -> Operands (Vec n single) -> CodeGen arch (Operands tuple)
vecUnpack vecr (OP_Vec vec) = go vecr n
where
go :: VecR n' single tuple' -> Int -> CodeGen arch (Operands tuple')
go (VecRnil _) 0 = return $ OP_Unit
go (VecRnil _) _ = internalError "index mismatch"
go (VecRsucc vecr') i = do
xs <- go vecr' (i - 1)
x <- instr' $ ExtractElement (fromIntegral i - 1) vec
return $ OP_Pair xs (ir singleTp x)
singleTp :: SingleType single -- GHC 8.4 cannot infer this type for some reason
VectorType n singleTp = vecRvector vecr
linearIndex :: ArrayVar aenv (Array sh e) -> Operands Int -> IROpenExp arch env aenv e
linearIndex (Var repr v) = linearIndexArray (irArray repr (aprj v aenv))
index :: ArrayVar aenv (Array sh e) -> Operands sh -> IROpenExp arch env aenv e
index (Var repr v) = indexArray (irArray repr (aprj v aenv))
shape :: ArrayVar aenv (Array sh e) -> Operands sh
shape (Var repr v) = irArrayShape (irArray repr (aprj v aenv))
pair :: Operands t1 -> Operands t2 -> IROpenExp arch env aenv (t1, t2)
pair a b = return $ OP_Pair a b
bool :: IROpenExp arch env aenv PrimBool
-> IROpenExp arch env aenv Bool
bool p = instr . IntToBool integralType . op integralType =<< p
primbool :: IROpenExp arch env aenv Bool
-> IROpenExp arch env aenv PrimBool
primbool b = instr . BoolToInt integralType . A.unbool =<< b
cond :: TypeR a
-> IROpenExp arch env aenv PrimBool
-> IROpenExp arch env aenv a
-> IROpenExp arch env aenv a
-> IROpenExp arch env aenv a
cond tp p t e =
A.ifThenElse (tp, bool p) t e
while :: TypeR a
-> IROpenFun1 arch env aenv (a -> PrimBool)
-> IROpenFun1 arch env aenv (a -> a)
-> IROpenExp arch env aenv a
-> IROpenExp arch env aenv a
while tp p f x =
L.while tp (bool . app1 p) (app1 f) =<< x
land :: Operands PrimBool
-> Operands PrimBool
-> IROpenExp arch env aenv PrimBool
land x y = do
x' <- instr (IntToBool integralType (op integralType x))
y' <- instr (IntToBool integralType (op integralType y))
primbool (A.land x' y')
lor :: Operands PrimBool
-> Operands PrimBool
-> IROpenExp arch env aenv PrimBool
lor x y = do
x' <- instr (IntToBool integralType (op integralType x))
y' <- instr (IntToBool integralType (op integralType y))
primbool (A.lor x' y')
foreignE :: A.Foreign asm
=> TypeR b
-> asm (a -> b)
-> Fun () (a -> b)
-> Operands a
-> IRExp arch () b
foreignE _ asm no x =
case foreignExp asm of
Just f -> app1 f x
Nothing | Lam lhs (Body b) <- no -> llvmOfOpenExp b (Empty `pushE` (lhs, x)) IM.empty
_ -> error "when a grid's misaligned with another behind / that's a moiré..."
coerce :: ScalarType a -> ScalarType b -> Operands a -> IROpenExp arch env aenv b
coerce s t x
| Just Refl <- matchScalarType s t = return $ x
| otherwise = ir t <$> instr' (BitCast t (op s x))
primFun :: PrimFun (a -> r)
-> OpenExp env aenv a
-> IROpenExp arch env aenv r
primFun f x =
case f of
PrimAdd t -> A.uncurry (A.add t) =<< cvtE x
PrimSub t -> A.uncurry (A.sub t) =<< cvtE x
PrimMul t -> A.uncurry (A.mul t) =<< cvtE x
PrimNeg t -> A.negate t =<< cvtE x
PrimAbs t -> A.abs t =<< cvtE x
PrimSig t -> A.signum t =<< cvtE x
PrimQuot t -> A.uncurry (A.quot t) =<< cvtE x
PrimRem t -> A.uncurry (A.rem t) =<< cvtE x
PrimQuotRem t -> A.uncurry (A.quotRem t) =<< cvtE x
PrimIDiv t -> A.uncurry (A.idiv t) =<< cvtE x
PrimMod t -> A.uncurry (A.mod t) =<< cvtE x
PrimDivMod t -> A.uncurry (A.divMod t) =<< cvtE x
PrimBAnd t -> A.uncurry (A.band t) =<< cvtE x
PrimBOr t -> A.uncurry (A.bor t) =<< cvtE x
PrimBXor t -> A.uncurry (A.xor t) =<< cvtE x
PrimBNot t -> A.complement t =<< cvtE x
PrimBShiftL t -> A.uncurry (A.shiftL t) =<< cvtE x
PrimBShiftR t -> A.uncurry (A.shiftR t) =<< cvtE x
PrimBRotateL t -> A.uncurry (A.rotateL t) =<< cvtE x
PrimBRotateR t -> A.uncurry (A.rotateR t) =<< cvtE x
PrimPopCount t -> A.popCount t =<< cvtE x
PrimCountLeadingZeros t -> A.countLeadingZeros t =<< cvtE x
PrimCountTrailingZeros t -> A.countTrailingZeros t =<< cvtE x
PrimFDiv t -> A.uncurry (A.fdiv t) =<< cvtE x
PrimRecip t -> A.recip t =<< cvtE x
PrimSin t -> A.sin t =<< cvtE x
PrimCos t -> A.cos t =<< cvtE x
PrimTan t -> A.tan t =<< cvtE x
PrimSinh t -> A.sinh t =<< cvtE x
PrimCosh t -> A.cosh t =<< cvtE x
PrimTanh t -> A.tanh t =<< cvtE x
PrimAsin t -> A.asin t =<< cvtE x
PrimAcos t -> A.acos t =<< cvtE x
PrimAtan t -> A.atan t =<< cvtE x
PrimAsinh t -> A.asinh t =<< cvtE x
PrimAcosh t -> A.acosh t =<< cvtE x
PrimAtanh t -> A.atanh t =<< cvtE x
PrimAtan2 t -> A.uncurry (A.atan2 t) =<< cvtE x
PrimExpFloating t -> A.exp t =<< cvtE x
PrimFPow t -> A.uncurry (A.fpow t) =<< cvtE x
PrimSqrt t -> A.sqrt t =<< cvtE x
PrimLog t -> A.log t =<< cvtE x
PrimLogBase t -> A.uncurry (A.logBase t) =<< cvtE x
PrimTruncate ta tb -> A.truncate ta tb =<< cvtE x
PrimRound ta tb -> A.round ta tb =<< cvtE x
PrimFloor ta tb -> A.floor ta tb =<< cvtE x
PrimCeiling ta tb -> A.ceiling ta tb =<< cvtE x
PrimMax t -> A.uncurry (A.max t) =<< cvtE x
PrimMin t -> A.uncurry (A.min t) =<< cvtE x
PrimFromIntegral ta tb -> A.fromIntegral ta tb =<< cvtE x
PrimToFloating ta tb -> A.toFloating ta tb =<< cvtE x
PrimLAnd -> A.uncurry land =<< cvtE x
PrimLOr -> A.uncurry lor =<< cvtE x
PrimIsNaN t -> primbool $ A.isNaN t =<< cvtE x
PrimIsInfinite t -> primbool $ A.isInfinite t =<< cvtE x
PrimLt t -> primbool $ A.uncurry (A.lt t) =<< cvtE x
PrimGt t -> primbool $ A.uncurry (A.gt t) =<< cvtE x
PrimLtEq t -> primbool $ A.uncurry (A.lte t) =<< cvtE x
PrimGtEq t -> primbool $ A.uncurry (A.gte t) =<< cvtE x
PrimEq t -> primbool $ A.uncurry (A.eq t) =<< cvtE x
PrimNEq t -> primbool $ A.uncurry (A.neq t) =<< cvtE x
PrimLNot -> primbool $ A.lnot =<< bool (cvtE x)
-- no missing patterns, whoo!
-- | Extract the head of an index
--
indexHead :: Operands (sh, sz) -> Operands sz
indexHead (OP_Pair _ sz) = sz
-- | Extract the tail of an index
--
indexTail :: Operands (sh, sz) -> Operands sh
indexTail (OP_Pair sh _) = sh
-- | Construct an index from the head and tail
--
indexCons :: Operands sh -> Operands sz -> Operands (sh, sz)
indexCons sh sz = OP_Pair sh sz
-- | Number of elements contained within a shape
--
shapeSize :: ShapeR sh -> Operands sh -> CodeGen arch (Operands Int)
shapeSize ShapeRz OP_Unit
= return $ A.liftInt 1
shapeSize (ShapeRsnoc shr) (OP_Pair sh sz)
= case shr of
ShapeRz -> return sz
_ -> do
a <- shapeSize shr sh
b <- A.mul numType a sz
return b
-- | Convert a multidimensional array index into a linear index
--
intOfIndex :: ShapeR sh -> Operands sh -> Operands sh -> CodeGen arch (Operands Int)
intOfIndex ShapeRz OP_Unit OP_Unit
= return $ A.liftInt 0
intOfIndex (ShapeRsnoc shr) (OP_Pair sh sz) (OP_Pair ix i)
-- If we short-circuit the last dimension, we can avoid inserting
-- a multiply by zero and add of the result.
= case shr of
ShapeRz -> return i
_ -> do
a <- intOfIndex shr sh ix
b <- A.mul numType a sz
c <- A.add numType b i
return c
-- | Convert a linear index into into a multidimensional index
--
indexOfInt :: ShapeR sh -> Operands sh -> Operands Int -> CodeGen arch (Operands sh)
indexOfInt ShapeRz OP_Unit _
= return OP_Unit
indexOfInt (ShapeRsnoc shr) (OP_Pair sh sz) i
= do
i' <- A.quot integralType i sz
-- If we assume the index is in range, there is no point computing
-- the remainder of the highest dimension since (i < sz) must hold
r <- case shr of
ShapeRz -> return i -- TODO: in debug mode assert (i < sz)
_ -> A.rem integralType i sz
sh' <- indexOfInt shr sh i'
return $ OP_Pair sh' r
-- | Read an element at a multidimensional index
--
indexArray :: IRArray (Array sh e) -> Operands sh -> IROpenExp arch env aenv e
indexArray arr ix = linearIndexArray arr =<< intOfIndex (arrayRshape $ irArrayRepr arr) (irArrayShape arr) ix
-- | Read an element at a linear index
--
linearIndexArray :: IRArray (Array sh e) -> Operands Int -> IROpenExp arch env aenv e
linearIndexArray = readArray TypeInt
pushE :: Val env -> (ELeftHandSide t env env', Operands t) -> Val env'
pushE env (LeftHandSideSingle _ , e) = env `Push` e
pushE env (LeftHandSideWildcard _, _) = env
pushE env (LeftHandSidePair l1 l2, (OP_Pair e1 e2)) = pushE env (l1, e1) `pushE` (l2, e2)