module Data.Array.Knead.Parameterized.PhysicalHull (
render,
MapFilter(..),
mapFilter,
FilterOuter(..),
filterOuter,
Scatter(..),
scatter,
ScatterMaybe(..),
scatterMaybe,
MapAccumLSimple(..),
mapAccumLSimple,
MapAccumLSequence(..),
mapAccumLSequence,
MapAccumL(..),
mapAccumL,
FoldOuterL(..),
foldOuterL,
AddDimension(..),
addDimension,
) where
import qualified Data.Array.Knead.Parameterized.Private as Sym
import qualified Data.Array.Knead.Simple.PhysicalPrivate as Priv
import qualified Data.Array.Knead.Simple.Private as Core
import qualified Data.Array.Knead.Shape as Shape
import qualified Data.Array.Knead.Expression as Expr
import Data.Array.Knead.Simple.PhysicalPrivate (MarshalPtr)
import Data.Array.Comfort.Storable.Unchecked (Array(Array))
import qualified LLVM.DSL.Execution as Code
import LLVM.DSL.Expression (Exp, unExp)
import qualified LLVM.Extra.Multi.Value as MultiValue
import qualified LLVM.Extra.Storable as Storable
import qualified LLVM.Extra.Marshal as Marshal
import qualified LLVM.Extra.Memory as Memory
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Core as LLVM
import Foreign.Marshal.Array (allocaArray, )
import Foreign.Marshal.Alloc (alloca, )
import Foreign.Storable (Storable, peek, peekElemOff, )
import Foreign.ForeignPtr (ForeignPtr, withForeignPtr, mallocForeignPtrArray, )
import Foreign.Ptr (FunPtr, Ptr, )
import Control.Exception (bracket, )
import Control.Monad.HT (void, )
import Control.Applicative (liftA2, )
mallocArray :: (Storable a) => Shape.Size -> IO (ForeignPtr a)
mallocArray = mallocForeignPtrArray . fromIntegral
type Importer f = FunPtr f -> f
foreign import ccall safe "dynamic" callShaper ::
Importer (LLVM.Ptr param -> LLVM.Ptr shape -> IO Shape.Size)
foreign import ccall safe "dynamic" callFill ::
Importer (LLVM.Ptr param -> LLVM.Ptr shape -> Ptr a -> IO ())
materialize ::
(Shape.C sh, Marshal.MV sh, Storable.C a) =>
String ->
(core -> Exp sh) ->
(core ->
LLVM.Value (MarshalPtr sh) -> LLVM.Value (Ptr a) ->
LLVM.CodeGenFunction () ()) ->
Sym.Hull p core -> IO (p -> IO (Array sh a))
materialize name shape fill (Sym.Hull core create delete) = do
(fsh, farr) <-
Code.compile name $
liftA2 (,)
(Code.createFunction callShaper "shape" $
\paramPtr resultPtr -> do
param <- Memory.load paramPtr
sh <- unExp $ shape $ core param
Memory.store sh resultPtr
Shape.size sh >>= LLVM.ret)
(Code.createFunction callFill "fill" $
\paramPtr shapePtr bufferPtr -> do
param <- Memory.load paramPtr
fill (core param) shapePtr bufferPtr
LLVM.ret ())
return $ \p ->
bracket (create p) (delete . fst) $ \(_ctx, param) ->
Marshal.alloca $ \shptr ->
Marshal.with param $ \paramPtr -> do
fptr <- mallocArray =<< fsh paramPtr shptr
withForeignPtr fptr $ farr paramPtr shptr
sh <- Marshal.peek shptr
return (Array sh fptr)
foreign import ccall safe "dynamic" callFillExpArray ::
Importer (LLVM.Ptr param -> Ptr final -> LLVM.Ptr shape -> Ptr a -> IO ())
materializeExpArray ::
(Shape.C sh, Marshal.MV sh, Storable.C a, Storable.C b) =>
String ->
(core -> Exp sh) ->
(core ->
LLVM.Value (Ptr b) ->
LLVM.Value (MarshalPtr sh) ->
LLVM.Value (Ptr a) ->
LLVM.CodeGenFunction () ()) ->
Sym.Hull p core -> IO (p -> IO (b, Array sh a))
materializeExpArray name shape fill (Sym.Hull core create delete) = do
(fsh, farr) <-
Code.compile name $
liftA2 (,)
(Code.createFunction callShaper "shape" $
\paramPtr resultPtr -> do
param <- Memory.load paramPtr
sh <- unExp $ shape $ core param
Memory.store sh resultPtr
Shape.size sh >>= LLVM.ret)
(Code.createFunction callFillExpArray "fill" $
\paramPtr finalPtr shapePtr bufferPtr -> do
param <- Memory.load paramPtr
fill (core param) finalPtr shapePtr bufferPtr
LLVM.ret ())
return $ \p ->
bracket (create p) (delete . fst) $ \(_ctx, param) ->
Marshal.alloca $ \shptr ->
alloca $ \finalPtr ->
Marshal.with param $ \paramPtr -> do
fptr <- mallocArray =<< fsh paramPtr shptr
withForeignPtr fptr $ farr paramPtr finalPtr shptr
sh <- Marshal.peek shptr
final <- peek finalPtr
return (final, Array sh fptr)
foreign import ccall safe "dynamic" callShaper2 ::
Importer
(LLVM.Ptr param ->
LLVM.Ptr shapeA -> LLVM.Ptr shapeB -> Ptr Shape.Size -> IO ())
foreign import ccall safe "dynamic" callFill2 ::
Importer
(LLVM.Ptr param ->
LLVM.Ptr shapeA -> Ptr a -> LLVM.Ptr shapeB -> Ptr b -> IO ())
materialize2 ::
(Shape.C sha, Marshal.MV sha,
Shape.C shb, Marshal.MV shb,
Storable.C a, Storable.C b) =>
String ->
(core -> Exp (sha,shb)) ->
(core ->
(LLVM.Value (MarshalPtr sha), LLVM.Value (Ptr a)) ->
(LLVM.Value (MarshalPtr shb), LLVM.Value (Ptr b)) ->
LLVM.CodeGenFunction () ()) ->
Sym.Hull p core -> IO (p -> IO (Array sha a, Array shb b))
materialize2 name shape fill (Sym.Hull core create delete) = do
(fsh, farr) <-
Code.compile name $
liftA2 (,)
(Code.createFunction callShaper2 "shape" $
\paramPtr shapeAPtr shapeBPtr sizesPtr -> do
param <- Memory.load paramPtr
(sha,shb) <- fmap MultiValue.unzip $ unExp $ shape $ core param
Memory.store sha shapeAPtr
Memory.store shb shapeBPtr
sizeAPtr <- LLVM.bitcast sizesPtr
flip LLVM.store sizeAPtr =<< Shape.size sha
sizeBPtr <- A.advanceArrayElementPtr sizeAPtr
flip LLVM.store sizeBPtr =<< Shape.size shb
LLVM.ret ())
(Code.createFunction callFill2 "fill" $
\paramPtr shapeAPtr bufferAPtr shapeBPtr bufferBPtr -> do
param <- Memory.load paramPtr
fill (core param) (shapeAPtr, bufferAPtr) (shapeBPtr, bufferBPtr)
LLVM.ret ())
return $ \p ->
bracket (create p) (delete . fst) $ \(_ctx, param) ->
Marshal.alloca $ \shaPtr ->
Marshal.alloca $ \shbPtr ->
allocaArray 2 $ \sizesPtr ->
Marshal.with param $ \paramPtr -> do
fsh paramPtr shaPtr shbPtr sizesPtr
afptr <- mallocArray =<< peekElemOff sizesPtr 0
bfptr <- mallocArray =<< peekElemOff sizesPtr 1
withForeignPtr afptr $ \aptr ->
withForeignPtr bfptr $ \bptr ->
farr paramPtr shaPtr aptr shbPtr bptr
sha <- Marshal.peek shaPtr
shb <- Marshal.peek shbPtr
return (Array sha afptr, Array shb bfptr)
render ::
(Shape.C sh, Shape.Index sh ~ ix, Marshal.MV sh,
Storable.C a) =>
Sym.Hull p (Core.Array sh a) -> IO (p -> IO (Array sh a))
render =
materialize "render" Core.shape
(\(Core.Array esh code) shapePtr bufferPtr -> do
let step ix p = flip Storable.storeNextMultiValue p =<< code ix
sh <- Shape.load esh shapePtr
void $ Shape.loop step sh bufferPtr)
data Scatter sh0 sh1 a =
Scatter {
scatterAccum :: Exp a -> Exp a -> Exp a,
scatterInit :: Core.Array sh1 a,
scatterMap :: Core.Array sh0 (Shape.Index sh1, a)
}
scatter ::
(Shape.C sh0, Shape.Index sh0 ~ ix0,
Shape.C sh1, Shape.Index sh1 ~ ix1, Marshal.MV sh1,
Storable.C a) =>
Sym.Hull p (Scatter sh0 sh1 a) -> IO (p -> IO (Array sh1 a))
scatter =
materialize "scatter"
(Core.shape . scatterInit)
(\(Scatter accum arrInit arrMap) ->
Priv.scatter accum arrInit arrMap)
data ScatterMaybe sh0 sh1 a =
ScatterMaybe {
scatterMaybeAccum :: Exp a -> Exp a -> Exp a,
scatterMaybeInit :: Core.Array sh1 a,
scatterMaybeMap :: Core.Array sh0 (Maybe (Shape.Index sh1, a))
}
scatterMaybe ::
(Shape.C sh0, Shape.Index sh0 ~ ix0,
Shape.C sh1, Shape.Index sh1 ~ ix1, Marshal.MV sh1,
Storable.C a) =>
Sym.Hull p (ScatterMaybe sh0 sh1 a) -> IO (p -> IO (Array sh1 a))
scatterMaybe =
materialize "scatterMaybe"
(Core.shape . scatterMaybeInit)
(\(ScatterMaybe accum arrInit arrMap) ->
Priv.scatterMaybe accum arrInit arrMap)
data MapAccumLSimple sh n acc a b =
MapAccumLSimple {
mapAccumLSimpleAccum :: Exp acc -> Exp a -> Exp (acc,b),
mapAccumLSimpleInit :: Core.Array sh acc,
mapAccumLSimpleArray :: Core.Array (sh, n) a
}
mapAccumLSimple ::
(Shape.C sh, Marshal.MV sh,
Shape.C n, Marshal.MV n,
MultiValue.C acc, Storable.C a, Storable.C b) =>
Sym.Hull p (MapAccumLSimple sh n acc a b) -> IO (p -> IO (Array (sh,n) b))
mapAccumLSimple =
materialize "mapAccumLSimple"
(Core.shape . mapAccumLSimpleArray)
(\(MapAccumLSimple f arrInit arrData) ->
Priv.mapAccumLSimple f arrInit arrData)
data MapAccumLSequence n acc final a b =
MapAccumLSequence {
mapAccumLSequenceAccum :: Exp acc -> Exp a -> Exp (acc,b),
mapAccumLSequenceFinal :: Exp acc -> Exp final,
mapAccumLSequenceInit :: Exp acc,
mapAccumLSequenceArray :: Core.Array n a
}
mapAccumLSequence ::
(Shape.C n, Marshal.MV n, MultiValue.C acc,
Storable.C final, MultiValue.C final,
Storable.C a, Storable.C b) =>
Sym.Hull p (MapAccumLSequence n acc final a b) ->
IO (p -> IO (final, Array n b))
mapAccumLSequence =
materializeExpArray "mapAccumLSequence"
(Core.shape . mapAccumLSequenceArray)
(\(MapAccumLSequence f final expInit arr) ->
Priv.mapAccumLSequence f final expInit arr)
data MapAccumL sh n acc final a b =
MapAccumL {
mapAccumLAccum :: Exp acc -> Exp a -> Exp (acc,b),
mapAccumLFinal :: Exp acc -> Exp final,
mapAccumLInit :: Core.Array sh acc,
mapAccumLArray :: Core.Array (sh, n) a
}
mapAccumL ::
(Shape.C sh, Marshal.MV sh,
Shape.C n, Marshal.MV n,
MultiValue.C acc,
Storable.C final, MultiValue.C final,
Storable.C a, Storable.C b) =>
Sym.Hull p (MapAccumL sh n acc final a b) ->
IO (p -> IO (Array sh final, Array (sh,n) b))
mapAccumL =
materialize2 "mapAccumL"
(\core ->
Expr.zip
(Core.shape $ mapAccumLInit core)
(Core.shape $ mapAccumLArray core))
(\(MapAccumL f final arrInit arrData) ->
Priv.mapAccumL f final arrInit arrData)
data FoldOuterL n sh a b =
FoldOuterL {
foldOuterLAccum :: Exp a -> Exp b -> Exp a,
foldOuterLInit :: Core.Array sh a,
foldOuterLArray :: Core.Array (n,sh) b
}
foldOuterL ::
(Shape.C n, Marshal.MV n,
Shape.C sh, Marshal.MV sh,
Storable.C a) =>
Sym.Hull p (FoldOuterL n sh a b) -> IO (p -> IO (Array sh a))
foldOuterL =
materialize "foldOuterL"
(Core.shape . foldOuterLInit)
(\(FoldOuterL f arrInit arrData) -> Priv.foldOuterL f arrInit arrData)
data MapFilter n a b =
MapFilter {
mapFilterMap :: Exp a -> Exp b,
mapFilterPredicate :: Exp a -> Exp Bool,
mapFilterArray :: Core.Array n a
}
mapFilter ::
(Shape.Sequence n, Marshal.MV n, Storable.C b) =>
Sym.Hull p (MapFilter n a b) -> IO (p -> IO (Array n b))
mapFilter =
materialize "mapFilter"
(Core.shape . mapFilterArray)
(\(MapFilter f p arr) shapePtr bufferPtr ->
flip Memory.store shapePtr
=<< Priv.mapFilter f p arr shapePtr bufferPtr)
data FilterOuter n sh a =
FilterOuter {
filterOuterPredicate :: Core.Array n Bool,
filterOuterArray :: Core.Array (n,sh) a
}
filterOuter ::
(Shape.Sequence n, Marshal.MV n,
Shape.C sh, Marshal.MV sh,
Storable.C a) =>
Sym.Hull p (FilterOuter n sh a) -> IO (p -> IO (Array (n,sh) a))
filterOuter =
materialize "filterOuter"
(Core.shape . filterOuterArray)
(\(FilterOuter p arr) shapePtr bufferPtr ->
flip Memory.store shapePtr
=<< Priv.filterOuter p arr shapePtr bufferPtr)
data AddDimension sh n a b =
AddDimension {
addDimensionSize :: Exp n,
addDimensionSelect :: Exp (Shape.Index n) -> Exp a -> Exp b,
addDimensionArray :: Core.Array sh a
}
addDimension ::
(Shape.C sh, Marshal.MV sh,
Shape.C n, Marshal.MV n,
Storable.C b) =>
Sym.Hull p (AddDimension sh n a b) -> IO (p -> IO (Array (sh,n) b))
addDimension =
materialize "addDimension"
(\r -> Expr.zip (Core.shape (addDimensionArray r)) (addDimensionSize r))
(\(AddDimension n select arr) -> Priv.addDimension n select arr)