{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ForeignFunctionInterface #-}
module Data.Array.Knead.Parameterized.Physical (
   Phys.Array,
   Array.shape,
   Phys.fromList,
   feed,
   the,
   theMarshal,
   render,
   renderShape,
   mapAccumLSimple,
   foldOuterL,
   scatter,
   scatterMaybe,
   permute,
   ) where

import qualified Data.Array.Knead.Parameterized.PhysicalHull as PhysHull
import qualified Data.Array.Knead.Parameterized.Private as Sym
import qualified Data.Array.Knead.Symbolic.Physical as Phys
import qualified Data.Array.Knead.Symbolic.Private as Core
import qualified Data.Array.Knead.Shape as Shape
import qualified Data.Array.Knead.Expression as Expr
import Data.Array.Knead.Code (getElementPtr)

import qualified Data.Array.Comfort.Storable.Unchecked as Array

import qualified LLVM.DSL.Parameter as Param
import qualified LLVM.DSL.Execution as Code
import LLVM.DSL.Expression (Exp, unExp)

import qualified LLVM.Extra.Multi.Value.Storable as Storable
import qualified LLVM.Extra.Multi.Value.Marshal as Marshal
import qualified LLVM.Extra.Multi.Value as MultiValue
import qualified LLVM.Extra.Memory as Memory

import qualified LLVM.Core as LLVM

import Foreign.Marshal.Alloc (alloca, )
import Foreign.Storable (peek, )
import Foreign.ForeignPtr (withForeignPtr, touchForeignPtr, )
import Foreign.Ptr (FunPtr, Ptr, )

import Control.Exception (bracket, )
import Control.Monad.HT ((<=<), )
import Control.Applicative (liftA2, )
import Data.Tuple.HT (mapFst, )


{-# INLINE feed #-}
feed ::
   (Shape.C sh, Marshal.C sh, Storable.C a) =>
   Param.T p (Phys.Array sh a) -> Sym.Array p sh a
feed :: forall sh a p.
(C sh, C sh, C a) =>
T p (Array sh a) -> Array p sh a
feed T p (Array sh a)
arr =
   T p sh
-> (forall parameters.
    C parameters =>
    (p -> parameters) -> (T parameters -> T sh) -> Array p sh a)
-> Array p sh a
forall b p a.
C b =>
T p b
-> (forall parameters.
    C parameters =>
    (p -> parameters) -> (T parameters -> T b) -> a)
-> a
Param.withMulti ((Array sh a -> sh) -> T p (Array sh a) -> T p sh
forall a b. (a -> b) -> T p a -> T p b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Array sh a -> sh
forall sh a. Array sh a -> sh
Array.shape T p (Array sh a)
arr) ((forall parameters.
  C parameters =>
  (p -> parameters) -> (T parameters -> T sh) -> Array p sh a)
 -> Array p sh a)
-> (forall parameters.
    C parameters =>
    (p -> parameters) -> (T parameters -> T sh) -> Array p sh a)
-> Array p sh a
forall a b. (a -> b) -> a -> b
$ \p -> parameters
getShape T parameters -> T sh
valueShape ->
   (T (parameters, Ptr a) -> Array sh a)
-> (p -> IO (ForeignPtr a, (parameters, Ptr a)))
-> (ForeignPtr a -> IO ())
-> Array p sh a
forall p sh a parameter context.
C parameter =>
(T parameter -> Array sh a)
-> (p -> IO (context, parameter))
-> (context -> IO ())
-> Array p sh a
Sym.Array
      (\T (parameters, Ptr a)
p ->
         case (T parameters -> T sh)
-> (T parameters, T (Ptr a)) -> (T sh, T (Ptr a))
forall a c b. (a -> c) -> (a, b) -> (c, b)
mapFst T parameters -> T sh
valueShape ((T parameters, T (Ptr a)) -> (T sh, T (Ptr a)))
-> (T parameters, T (Ptr a)) -> (T sh, T (Ptr a))
forall a b. (a -> b) -> a -> b
$ T (parameters, Ptr a) -> (T parameters, T (Ptr a))
forall a b. T (a, b) -> (T a, T b)
MultiValue.unzip T (parameters, Ptr a)
p of
            (T sh
sh, MultiValue.Cons Repr (Ptr a)
ptr) ->
               Exp sh -> (forall {r}. Val (Index sh) -> Code r a) -> Array sh a
forall sh a.
Exp sh -> (forall r. Val (Index sh) -> Code r a) -> Array sh a
Core.Array (T sh -> Exp sh
forall a. T a -> Exp a
forall (val :: * -> *) a. Value val => T a -> val a
Expr.lift0 T sh
sh) ((forall {r}. Val (Index sh) -> Code r a) -> Array sh a)
-> (forall {r}. Val (Index sh) -> Code r a) -> Array sh a
forall a b. (a -> b) -> a -> b
$
                  Value (Ptr a) -> CodeGenFunction r (Val a)
forall a r. C a => Value (Ptr a) -> CodeGenFunction r (T a)
forall r. Value (Ptr a) -> CodeGenFunction r (Val a)
Storable.load (Value (Ptr a) -> CodeGenFunction r (Val a))
-> (Val (Index sh) -> CodeGenFunction r (Value (Ptr a)))
-> Val (Index sh)
-> CodeGenFunction r (Val a)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< T sh
-> Value (Ptr a)
-> Val (Index sh)
-> CodeGenFunction r (Value (Ptr a))
forall sh ix a r.
(C sh, Index sh ~ ix, C a) =>
T sh -> Value (Ptr a) -> T ix -> CodeGenFunction r (Value (Ptr a))
getElementPtr T sh
sh Repr (Ptr a)
Value (Ptr a)
ptr)
      (\p
p ->
         case Array sh a -> ForeignPtr a
forall sh a. Array sh a -> ForeignPtr a
Array.buffer (Array sh a -> ForeignPtr a) -> Array sh a -> ForeignPtr a
forall a b. (a -> b) -> a -> b
$ T p (Array sh a) -> p -> Array sh a
forall p a. T p a -> p -> a
Param.get T p (Array sh a)
arr p
p of
            ForeignPtr a
fptr ->
               ForeignPtr a
-> (Ptr a -> IO (ForeignPtr a, (parameters, Ptr a)))
-> IO (ForeignPtr a, (parameters, Ptr a))
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
fptr ((Ptr a -> IO (ForeignPtr a, (parameters, Ptr a)))
 -> IO (ForeignPtr a, (parameters, Ptr a)))
-> (Ptr a -> IO (ForeignPtr a, (parameters, Ptr a)))
-> IO (ForeignPtr a, (parameters, Ptr a))
forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr ->
                  (ForeignPtr a, (parameters, Ptr a))
-> IO (ForeignPtr a, (parameters, Ptr a))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ForeignPtr a
fptr, (p -> parameters
getShape p
p, Ptr a
ptr)))
      ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr


type Importer f = FunPtr f -> f

foreign import ccall safe "dynamic" callThe ::
   Importer (LLVM.Ptr param -> Ptr a -> IO ())


the ::
   (Shape.Scalar z, Storable.C a, MultiValue.C a) =>
   Sym.Array p z a -> IO (p -> IO a)
the :: forall z a p. (Scalar z, C a, C a) => Array p z a -> IO (p -> IO a)
the (Sym.Array T parameter -> Array z a
arr p -> IO (context, parameter)
create context -> IO ()
delete) = do
   Ptr (Struct (Repr parameter)) -> Ptr a -> IO ()
func <-
      String
-> Exec (Ptr (Struct (Repr parameter)) -> Ptr a -> IO ())
-> IO (Ptr (Struct (Repr parameter)) -> Ptr a -> IO ())
forall funcs. String -> Exec funcs -> IO funcs
Code.compile String
"the" (Exec (Ptr (Struct (Repr parameter)) -> Ptr a -> IO ())
 -> IO (Ptr (Struct (Repr parameter)) -> Ptr a -> IO ()))
-> Exec (Ptr (Struct (Repr parameter)) -> Ptr a -> IO ())
-> IO (Ptr (Struct (Repr parameter)) -> Ptr a -> IO ())
forall a b. (a -> b) -> a -> b
$
      Importer (Ptr (Struct (Repr parameter)) -> Ptr a -> IO ())
-> String
-> CodeGen (Ptr (Struct (Repr parameter)) -> Ptr a -> IO ())
-> Exec (Ptr (Struct (Repr parameter)) -> Ptr a -> IO ())
forall f.
(ExecutionFunction f, C f) =>
Importer f -> String -> CodeGen f -> Exec f
Code.createFunction Importer (Ptr (Struct (Repr parameter)) -> Ptr a -> IO ())
forall param a. Importer (Ptr param -> Ptr a -> IO ())
callThe String
"eval" (CodeGen (Ptr (Struct (Repr parameter)) -> Ptr a -> IO ())
 -> Exec (Ptr (Struct (Repr parameter)) -> Ptr a -> IO ()))
-> CodeGen (Ptr (Struct (Repr parameter)) -> Ptr a -> IO ())
-> Exec (Ptr (Struct (Repr parameter)) -> Ptr a -> IO ())
forall a b. (a -> b) -> a -> b
$
      \Value (Ptr (Struct (Repr parameter)))
paramPtr Value (Ptr a)
resultPtr -> do
         T parameter
param <- Value (Ptr (Struct (T parameter)))
-> CodeGenFunction () (T parameter)
forall llvmValue r.
C llvmValue =>
Value (Ptr (Struct llvmValue)) -> CodeGenFunction r llvmValue
forall r.
Value (Ptr (Struct (T parameter)))
-> CodeGenFunction r (T parameter)
Memory.load Value (Ptr (Struct (Repr parameter)))
Value (Ptr (Struct (T parameter)))
paramPtr
         case T parameter -> Array z a
arr T parameter
param of
            Core.Array Exp z
z forall r. Val (Index z) -> Code r a
code ->
               Val (Index z) -> Code () a
forall r. Val (Index z) -> Code r a
code (Exp z -> Val (Index z)
forall sh (val :: * -> *) (f :: * -> *).
(Scalar sh, Value val) =>
f sh -> val (Index sh)
forall (val :: * -> *) (f :: * -> *).
Value val =>
f z -> val (Index z)
Shape.zeroIndex Exp z
z) Code () a
-> (Val a -> CodeGenFunction () ()) -> CodeGenFunction () ()
forall a b.
CodeGenFunction () a
-> (a -> CodeGenFunction () b) -> CodeGenFunction () b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
               (Val a -> Value (Ptr a) -> CodeGenFunction () ())
-> Value (Ptr a) -> Val a -> CodeGenFunction () ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip Val a -> Value (Ptr a) -> CodeGenFunction () ()
forall r. Val a -> Value (Ptr a) -> CodeGenFunction r ()
forall a r. C a => T a -> Value (Ptr a) -> CodeGenFunction r ()
Storable.store Value (Ptr a)
resultPtr
   (p -> IO a) -> IO (p -> IO a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((p -> IO a) -> IO (p -> IO a)) -> (p -> IO a) -> IO (p -> IO a)
forall a b. (a -> b) -> a -> b
$ \p
p ->
      IO (context, parameter)
-> ((context, parameter) -> IO ())
-> ((context, parameter) -> IO a)
-> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (p -> IO (context, parameter)
create p
p) (context -> IO ()
delete (context -> IO ())
-> ((context, parameter) -> context)
-> (context, parameter)
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (context, parameter) -> context
forall a b. (a, b) -> a
fst) (((context, parameter) -> IO a) -> IO a)
-> ((context, parameter) -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \(context
_ctx, parameter
param) ->
      parameter -> (Ptr (Struct (Repr parameter)) -> IO a) -> IO a
forall a b. C a => a -> (Ptr (Struct a) -> IO b) -> IO b
Marshal.with parameter
param ((Ptr (Struct (Repr parameter)) -> IO a) -> IO a)
-> (Ptr (Struct (Repr parameter)) -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Ptr (Struct (Repr parameter))
pptr ->
      (Ptr a -> IO a) -> IO a
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca ((Ptr a -> IO a) -> IO a) -> (Ptr a -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Ptr a
aptr -> Ptr (Struct (Repr parameter)) -> Ptr a -> IO ()
func Ptr (Struct (Repr parameter))
pptr Ptr a
aptr IO () -> IO a -> IO a
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
aptr

foreign import ccall safe "dynamic" callTheMarshal ::
   Importer (LLVM.Ptr param -> LLVM.Ptr a -> IO ())

theMarshal ::
   (Shape.Scalar z, Marshal.C a) =>
   Sym.Array p z a -> IO (p -> IO a)
theMarshal :: forall z a p. (Scalar z, C a) => Array p z a -> IO (p -> IO a)
theMarshal (Sym.Array T parameter -> Array z a
arr p -> IO (context, parameter)
create context -> IO ()
delete) = do
   Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ()
func <-
      String
-> Exec
     (Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ())
-> IO
     (Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ())
forall funcs. String -> Exec funcs -> IO funcs
Code.compile String
"the-marshal" (Exec
   (Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ())
 -> IO
      (Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ()))
-> Exec
     (Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ())
-> IO
     (Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ())
forall a b. (a -> b) -> a -> b
$
      Importer
  (Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ())
-> String
-> CodeGen
     (Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ())
-> Exec
     (Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ())
forall f.
(ExecutionFunction f, C f) =>
Importer f -> String -> CodeGen f -> Exec f
Code.createFunction Importer
  (Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ())
forall param a. Importer (Ptr param -> Ptr a -> IO ())
callTheMarshal String
"eval" (CodeGen
   (Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ())
 -> Exec
      (Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ()))
-> CodeGen
     (Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ())
-> Exec
     (Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ())
forall a b. (a -> b) -> a -> b
$
      \Value (Ptr (Struct (Repr parameter)))
paramPtr Value (Ptr (Struct (Repr a)))
resultPtr -> do
         T parameter
param <- Value (Ptr (Struct (T parameter)))
-> CodeGenFunction () (T parameter)
forall llvmValue r.
C llvmValue =>
Value (Ptr (Struct llvmValue)) -> CodeGenFunction r llvmValue
forall r.
Value (Ptr (Struct (T parameter)))
-> CodeGenFunction r (T parameter)
Memory.load Value (Ptr (Struct (Repr parameter)))
Value (Ptr (Struct (T parameter)))
paramPtr
         case T parameter -> Array z a
arr T parameter
param of
            Core.Array Exp z
z forall r. Val (Index z) -> Code r a
code ->
               Val (Index z) -> Code () a
forall r. Val (Index z) -> Code r a
code (Exp z -> Val (Index z)
forall sh (val :: * -> *) (f :: * -> *).
(Scalar sh, Value val) =>
f sh -> val (Index sh)
forall (val :: * -> *) (f :: * -> *).
Value val =>
f z -> val (Index z)
Shape.zeroIndex Exp z
z) Code () a
-> (Val a -> CodeGenFunction () ()) -> CodeGenFunction () ()
forall a b.
CodeGenFunction () a
-> (a -> CodeGenFunction () b) -> CodeGenFunction () b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
               (Val a -> Value (Ptr (Struct (Repr a))) -> CodeGenFunction () ())
-> Value (Ptr (Struct (Repr a))) -> Val a -> CodeGenFunction () ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip Val a -> Value (Ptr (Struct (Repr a))) -> CodeGenFunction () ()
Val a -> Value (Ptr (Struct (Val a))) -> CodeGenFunction () ()
forall r.
Val a -> Value (Ptr (Struct (Val a))) -> CodeGenFunction r ()
forall llvmValue r.
C llvmValue =>
llvmValue -> Value (Ptr (Struct llvmValue)) -> CodeGenFunction r ()
Memory.store Value (Ptr (Struct (Repr a)))
resultPtr
   (p -> IO a) -> IO (p -> IO a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((p -> IO a) -> IO (p -> IO a)) -> (p -> IO a) -> IO (p -> IO a)
forall a b. (a -> b) -> a -> b
$ \p
p ->
      IO (context, parameter)
-> ((context, parameter) -> IO ())
-> ((context, parameter) -> IO a)
-> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (p -> IO (context, parameter)
create p
p) (context -> IO ()
delete (context -> IO ())
-> ((context, parameter) -> context)
-> (context, parameter)
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (context, parameter) -> context
forall a b. (a, b) -> a
fst) (((context, parameter) -> IO a) -> IO a)
-> ((context, parameter) -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \(context
_ctx, parameter
param) ->
      parameter -> (Ptr (Struct (Repr parameter)) -> IO a) -> IO a
forall a b. C a => a -> (Ptr (Struct a) -> IO b) -> IO b
Marshal.with parameter
param ((Ptr (Struct (Repr parameter)) -> IO a) -> IO a)
-> (Ptr (Struct (Repr parameter)) -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Ptr (Struct (Repr parameter))
pptr ->
      (Ptr (Struct (Repr a)) -> IO a) -> IO a
forall a b. IsType a => (Ptr a -> IO b) -> IO b
Marshal.alloca ((Ptr (Struct (Repr a)) -> IO a) -> IO a)
-> (Ptr (Struct (Repr a)) -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Ptr (Struct (Repr a))
aptr ->
         Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr a)) -> IO ()
func Ptr (Struct (Repr parameter))
pptr Ptr (Struct (Repr a))
aptr IO () -> IO a -> IO a
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
         Ptr (Struct (Repr a)) -> IO a
forall a struct.
(C a, Struct a ~ struct, Marshal struct) =>
Ptr struct -> IO a
Marshal.peek Ptr (Struct (Repr a))
aptr


foreign import ccall safe "dynamic" callShaper ::
   Importer (LLVM.Ptr param -> LLVM.Ptr shape -> IO Shape.Size)


renderShape ::
   (Shape.C sh, Marshal.C sh,
    Storable.C a, MultiValue.C a) =>
   Sym.Array p sh a -> IO (p -> IO (sh, Shape.Size))
renderShape :: forall sh a p.
(C sh, C sh, C a, C a) =>
Array p sh a -> IO (p -> IO (sh, Size))
renderShape (Sym.Array T parameter -> Array sh a
arr p -> IO (context, parameter)
create context -> IO ()
delete) = do
   Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr sh)) -> IO Size
fsh <-
      String
-> Exec
     (Ptr (Struct (Repr parameter))
      -> Ptr (Struct (Repr sh)) -> IO Size)
-> IO
     (Ptr (Struct (Repr parameter))
      -> Ptr (Struct (Repr sh)) -> IO Size)
forall funcs. String -> Exec funcs -> IO funcs
Code.compile String
"renderShape" (Exec
   (Ptr (Struct (Repr parameter))
    -> Ptr (Struct (Repr sh)) -> IO Size)
 -> IO
      (Ptr (Struct (Repr parameter))
       -> Ptr (Struct (Repr sh)) -> IO Size))
-> Exec
     (Ptr (Struct (Repr parameter))
      -> Ptr (Struct (Repr sh)) -> IO Size)
-> IO
     (Ptr (Struct (Repr parameter))
      -> Ptr (Struct (Repr sh)) -> IO Size)
forall a b. (a -> b) -> a -> b
$
      Importer
  (Ptr (Struct (Repr parameter))
   -> Ptr (Struct (Repr sh)) -> IO Size)
-> String
-> CodeGen
     (Ptr (Struct (Repr parameter))
      -> Ptr (Struct (Repr sh)) -> IO Size)
-> Exec
     (Ptr (Struct (Repr parameter))
      -> Ptr (Struct (Repr sh)) -> IO Size)
forall f.
(ExecutionFunction f, C f) =>
Importer f -> String -> CodeGen f -> Exec f
Code.createFunction Importer
  (Ptr (Struct (Repr parameter))
   -> Ptr (Struct (Repr sh)) -> IO Size)
forall param shape. Importer (Ptr param -> Ptr shape -> IO Size)
callShaper String
"shape" (CodeGen
   (Ptr (Struct (Repr parameter))
    -> Ptr (Struct (Repr sh)) -> IO Size)
 -> Exec
      (Ptr (Struct (Repr parameter))
       -> Ptr (Struct (Repr sh)) -> IO Size))
-> CodeGen
     (Ptr (Struct (Repr parameter))
      -> Ptr (Struct (Repr sh)) -> IO Size)
-> Exec
     (Ptr (Struct (Repr parameter))
      -> Ptr (Struct (Repr sh)) -> IO Size)
forall a b. (a -> b) -> a -> b
$
      \Value (Ptr (Struct (Repr parameter)))
paramPtr Value (Ptr (Struct (Repr sh)))
resultPtr -> do
        T parameter
param <- Value (Ptr (Struct (T parameter)))
-> CodeGenFunction Size (T parameter)
forall llvmValue r.
C llvmValue =>
Value (Ptr (Struct llvmValue)) -> CodeGenFunction r llvmValue
forall r.
Value (Ptr (Struct (T parameter)))
-> CodeGenFunction r (T parameter)
Memory.load Value (Ptr (Struct (Repr parameter)))
Value (Ptr (Struct (T parameter)))
paramPtr
        case T parameter -> Array sh a
arr T parameter
param of
           Core.Array Exp sh
esh forall r. Val (Index sh) -> Code r a
_code -> do
              T sh
sh <- Exp sh -> forall r. CodeGenFunction r (T sh)
forall a. Exp a -> forall r. CodeGenFunction r (T a)
unExp Exp sh
esh
              T sh -> Value (Ptr (Struct (T sh))) -> CodeGenFunction Size ()
forall r.
T sh -> Value (Ptr (Struct (T sh))) -> CodeGenFunction r ()
forall llvmValue r.
C llvmValue =>
llvmValue -> Value (Ptr (Struct llvmValue)) -> CodeGenFunction r ()
Memory.store T sh
sh Value (Ptr (Struct (Repr sh)))
Value (Ptr (Struct (T sh)))
resultPtr
              T sh -> CodeGenFunction Size (Value Size)
forall r. T sh -> CodeGenFunction r (Value Size)
forall sh r. C sh => T sh -> CodeGenFunction r (Value Size)
Shape.size T sh
sh
   (p -> IO (sh, Size)) -> IO (p -> IO (sh, Size))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((p -> IO (sh, Size)) -> IO (p -> IO (sh, Size)))
-> (p -> IO (sh, Size)) -> IO (p -> IO (sh, Size))
forall a b. (a -> b) -> a -> b
$ \p
p ->
      IO (context, parameter)
-> ((context, parameter) -> IO ())
-> ((context, parameter) -> IO (sh, Size))
-> IO (sh, Size)
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (p -> IO (context, parameter)
create p
p) (context -> IO ()
delete (context -> IO ())
-> ((context, parameter) -> context)
-> (context, parameter)
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (context, parameter) -> context
forall a b. (a, b) -> a
fst) (((context, parameter) -> IO (sh, Size)) -> IO (sh, Size))
-> ((context, parameter) -> IO (sh, Size)) -> IO (sh, Size)
forall a b. (a -> b) -> a -> b
$ \(context
_ctx, parameter
param) ->
      (Ptr (Struct (Repr sh)) -> IO (sh, Size)) -> IO (sh, Size)
forall a b. IsType a => (Ptr a -> IO b) -> IO b
Marshal.alloca ((Ptr (Struct (Repr sh)) -> IO (sh, Size)) -> IO (sh, Size))
-> (Ptr (Struct (Repr sh)) -> IO (sh, Size)) -> IO (sh, Size)
forall a b. (a -> b) -> a -> b
$ \Ptr (Struct (Repr sh))
shptr ->
      parameter
-> (Ptr (Struct (Repr parameter)) -> IO (sh, Size))
-> IO (sh, Size)
forall a b. C a => a -> (Ptr (Struct a) -> IO b) -> IO b
Marshal.with parameter
param ((Ptr (Struct (Repr parameter)) -> IO (sh, Size)) -> IO (sh, Size))
-> (Ptr (Struct (Repr parameter)) -> IO (sh, Size))
-> IO (sh, Size)
forall a b. (a -> b) -> a -> b
$ \Ptr (Struct (Repr parameter))
pptr -> do
         Size
n <- Ptr (Struct (Repr parameter)) -> Ptr (Struct (Repr sh)) -> IO Size
fsh Ptr (Struct (Repr parameter))
pptr Ptr (Struct (Repr sh))
shptr
         sh
sh <- Ptr (Struct (Repr sh)) -> IO sh
forall a struct.
(C a, Struct a ~ struct, Marshal struct) =>
Ptr struct -> IO a
Marshal.peek Ptr (Struct (Repr sh))
shptr
         (sh, Size) -> IO (sh, Size)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (sh
sh, Size
n)


render ::
   (Shape.C sh, Marshal.C sh, Storable.C a) =>
   Sym.Array p sh a -> IO (p -> IO (Phys.Array sh a))
render :: forall sh a p.
(C sh, C sh, C a) =>
Array p sh a -> IO (p -> IO (Array sh a))
render = Hull p (Array sh a) -> IO (p -> IO (Array sh a))
forall sh ix a p.
(C sh, Index sh ~ ix, C sh, C a) =>
Hull p (Array sh a) -> IO (p -> IO (Array sh a))
PhysHull.render (Hull p (Array sh a) -> IO (p -> IO (Array sh a)))
-> (Array p sh a -> Hull p (Array sh a))
-> Array p sh a
-> IO (p -> IO (Array sh a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array p sh a -> Hull p (Array sh a)
forall p sh a. Array p sh a -> Hull p (Array sh a)
Sym.arrayHull


mapAccumLSimple ::
   (Shape.C sh, Marshal.C sh,
    Shape.C n, Marshal.C n,
    MultiValue.C acc,
    Storable.C a, MultiValue.C a,
    Storable.C b, MultiValue.C b) =>
   (Exp acc -> Exp a -> Exp (acc,b)) ->
   Sym.Array p sh acc ->
   Sym.Array p (sh, n) a ->
   IO (p -> IO (Phys.Array (sh,n) b))
mapAccumLSimple :: forall sh n acc a b p.
(C sh, C sh, C n, C n, C acc, C a, C a, C b, C b) =>
(Exp acc -> Exp a -> Exp (acc, b))
-> Array p sh acc
-> Array p (sh, n) a
-> IO (p -> IO (Array (sh, n) b))
mapAccumLSimple Exp acc -> Exp a -> Exp (acc, b)
f Array p sh acc
arrInit Array p (sh, n) a
arrMap =
   Hull p (MapAccumLSimple sh n acc a b)
-> IO (p -> IO (Array (sh, n) b))
forall sh n acc a b p.
(C sh, C sh, C n, C n, C acc, C a, C b) =>
Hull p (MapAccumLSimple sh n acc a b)
-> IO (p -> IO (Array (sh, n) b))
PhysHull.mapAccumLSimple (Hull p (MapAccumLSimple sh n acc a b)
 -> IO (p -> IO (Array (sh, n) b)))
-> Hull p (MapAccumLSimple sh n acc a b)
-> IO (p -> IO (Array (sh, n) b))
forall a b. (a -> b) -> a -> b
$
      (Array sh acc -> Array (sh, n) a -> MapAccumLSimple sh n acc a b)
-> Hull p (Array sh acc)
-> Hull p (Array (sh, n) a)
-> Hull p (MapAccumLSimple sh n acc a b)
forall a b c. (a -> b -> c) -> Hull p a -> Hull p b -> Hull p c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 ((Exp acc -> Exp a -> Exp (acc, b))
-> Array sh acc -> Array (sh, n) a -> MapAccumLSimple sh n acc a b
forall sh n acc a b.
(Exp acc -> Exp a -> Exp (acc, b))
-> Array sh acc -> Array (sh, n) a -> MapAccumLSimple sh n acc a b
PhysHull.MapAccumLSimple Exp acc -> Exp a -> Exp (acc, b)
f)
         (Array p sh acc -> Hull p (Array sh acc)
forall p sh a. Array p sh a -> Hull p (Array sh a)
Sym.arrayHull Array p sh acc
arrInit)
         (Array p (sh, n) a -> Hull p (Array (sh, n) a)
forall p sh a. Array p sh a -> Hull p (Array sh a)
Sym.arrayHull Array p (sh, n) a
arrMap)

foldOuterL ::
   (Shape.C sh, Marshal.C sh,
    Shape.C n, Marshal.C n,
    Storable.C a, MultiValue.C a) =>
   (Exp a -> Exp b -> Exp a) ->
   Sym.Array p sh a ->
   Sym.Array p (n,sh) b ->
   IO (p -> IO (Phys.Array sh a))
foldOuterL :: forall sh n a b p.
(C sh, C sh, C n, C n, C a, C a) =>
(Exp a -> Exp b -> Exp a)
-> Array p sh a -> Array p (n, sh) b -> IO (p -> IO (Array sh a))
foldOuterL Exp a -> Exp b -> Exp a
f Array p sh a
arrInit Array p (n, sh) b
arrMap =
   Hull p (FoldOuterL n sh a b) -> IO (p -> IO (Array sh a))
forall n sh a p b.
(C n, C n, C sh, C sh, C a) =>
Hull p (FoldOuterL n sh a b) -> IO (p -> IO (Array sh a))
PhysHull.foldOuterL (Hull p (FoldOuterL n sh a b) -> IO (p -> IO (Array sh a)))
-> Hull p (FoldOuterL n sh a b) -> IO (p -> IO (Array sh a))
forall a b. (a -> b) -> a -> b
$
      (Array sh a -> Array (n, sh) b -> FoldOuterL n sh a b)
-> Hull p (Array sh a)
-> Hull p (Array (n, sh) b)
-> Hull p (FoldOuterL n sh a b)
forall a b c. (a -> b -> c) -> Hull p a -> Hull p b -> Hull p c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 ((Exp a -> Exp b -> Exp a)
-> Array sh a -> Array (n, sh) b -> FoldOuterL n sh a b
forall n sh a b.
(Exp a -> Exp b -> Exp a)
-> Array sh a -> Array (n, sh) b -> FoldOuterL n sh a b
PhysHull.FoldOuterL Exp a -> Exp b -> Exp a
f)
         (Array p sh a -> Hull p (Array sh a)
forall p sh a. Array p sh a -> Hull p (Array sh a)
Sym.arrayHull Array p sh a
arrInit)
         (Array p (n, sh) b -> Hull p (Array (n, sh) b)
forall p sh a. Array p sh a -> Hull p (Array sh a)
Sym.arrayHull Array p (n, sh) b
arrMap)

scatter ::
   (Shape.C sh0, Shape.Index sh0 ~ ix0,
    Shape.C sh1, Shape.Index sh1 ~ ix1, Marshal.C sh1,
    Storable.C a, MultiValue.C a) =>
   (Exp a -> Exp a -> Exp a) ->
   Sym.Array p sh1 a ->
   Sym.Array p sh0 (ix1, a) -> IO (p -> IO (Phys.Array sh1 a))
scatter :: forall sh0 ix0 sh1 ix1 a p.
(C sh0, Index sh0 ~ ix0, C sh1, Index sh1 ~ ix1, C sh1, C a,
 C a) =>
(Exp a -> Exp a -> Exp a)
-> Array p sh1 a
-> Array p sh0 (ix1, a)
-> IO (p -> IO (Array sh1 a))
scatter Exp a -> Exp a -> Exp a
accum Array p sh1 a
arrBase Array p sh0 (ix1, a)
arrMap =
   Hull p (Scatter sh0 sh1 a) -> IO (p -> IO (Array sh1 a))
forall sh0 ix0 sh1 ix1 a p.
(C sh0, Index sh0 ~ ix0, C sh1, Index sh1 ~ ix1, C sh1, C a) =>
Hull p (Scatter sh0 sh1 a) -> IO (p -> IO (Array sh1 a))
PhysHull.scatter (Hull p (Scatter sh0 sh1 a) -> IO (p -> IO (Array sh1 a)))
-> Hull p (Scatter sh0 sh1 a) -> IO (p -> IO (Array sh1 a))
forall a b. (a -> b) -> a -> b
$
      (Array sh1 a -> Array sh0 (ix1, a) -> Scatter sh0 sh1 a)
-> Hull p (Array sh1 a)
-> Hull p (Array sh0 (ix1, a))
-> Hull p (Scatter sh0 sh1 a)
forall a b c. (a -> b -> c) -> Hull p a -> Hull p b -> Hull p c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 ((Exp a -> Exp a -> Exp a)
-> Array sh1 a -> Array sh0 (Index sh1, a) -> Scatter sh0 sh1 a
forall sh0 sh1 a.
(Exp a -> Exp a -> Exp a)
-> Array sh1 a -> Array sh0 (Index sh1, a) -> Scatter sh0 sh1 a
PhysHull.Scatter Exp a -> Exp a -> Exp a
accum)
         (Array p sh1 a -> Hull p (Array sh1 a)
forall p sh a. Array p sh a -> Hull p (Array sh a)
Sym.arrayHull Array p sh1 a
arrBase)
         (Array p sh0 (ix1, a) -> Hull p (Array sh0 (ix1, a))
forall p sh a. Array p sh a -> Hull p (Array sh a)
Sym.arrayHull Array p sh0 (ix1, a)
arrMap)

scatterMaybe ::
   (Shape.C sh0, Shape.Index sh0 ~ ix0,
    Shape.C sh1, Shape.Index sh1 ~ ix1, Marshal.C sh1,
    Storable.C a, MultiValue.C a) =>
   (Exp a -> Exp a -> Exp a) ->
   Sym.Array p sh1 a ->
   Sym.Array p sh0 (Maybe (ix1, a)) -> IO (p -> IO (Phys.Array sh1 a))
scatterMaybe :: forall sh0 ix0 sh1 ix1 a p.
(C sh0, Index sh0 ~ ix0, C sh1, Index sh1 ~ ix1, C sh1, C a,
 C a) =>
(Exp a -> Exp a -> Exp a)
-> Array p sh1 a
-> Array p sh0 (Maybe (ix1, a))
-> IO (p -> IO (Array sh1 a))
scatterMaybe Exp a -> Exp a -> Exp a
accum Array p sh1 a
arrBase Array p sh0 (Maybe (ix1, a))
arrMap =
   Hull p (ScatterMaybe sh0 sh1 a) -> IO (p -> IO (Array sh1 a))
forall sh0 ix0 sh1 ix1 a p.
(C sh0, Index sh0 ~ ix0, C sh1, Index sh1 ~ ix1, C sh1, C a) =>
Hull p (ScatterMaybe sh0 sh1 a) -> IO (p -> IO (Array sh1 a))
PhysHull.scatterMaybe (Hull p (ScatterMaybe sh0 sh1 a) -> IO (p -> IO (Array sh1 a)))
-> Hull p (ScatterMaybe sh0 sh1 a) -> IO (p -> IO (Array sh1 a))
forall a b. (a -> b) -> a -> b
$
      (Array sh1 a
 -> Array sh0 (Maybe (ix1, a)) -> ScatterMaybe sh0 sh1 a)
-> Hull p (Array sh1 a)
-> Hull p (Array sh0 (Maybe (ix1, a)))
-> Hull p (ScatterMaybe sh0 sh1 a)
forall a b c. (a -> b -> c) -> Hull p a -> Hull p b -> Hull p c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 ((Exp a -> Exp a -> Exp a)
-> Array sh1 a
-> Array sh0 (Maybe (Index sh1, a))
-> ScatterMaybe sh0 sh1 a
forall sh0 sh1 a.
(Exp a -> Exp a -> Exp a)
-> Array sh1 a
-> Array sh0 (Maybe (Index sh1, a))
-> ScatterMaybe sh0 sh1 a
PhysHull.ScatterMaybe Exp a -> Exp a -> Exp a
accum)
         (Array p sh1 a -> Hull p (Array sh1 a)
forall p sh a. Array p sh a -> Hull p (Array sh a)
Sym.arrayHull Array p sh1 a
arrBase)
         (Array p sh0 (Maybe (ix1, a)) -> Hull p (Array sh0 (Maybe (ix1, a)))
forall p sh a. Array p sh a -> Hull p (Array sh a)
Sym.arrayHull Array p sh0 (Maybe (ix1, a))
arrMap)

permute ::
   (Shape.C sh0, Shape.Index sh0 ~ ix0,
    Shape.C sh1, Shape.Index sh1 ~ ix1, Marshal.C sh1,
    Storable.C a, MultiValue.C a) =>
   (Exp a -> Exp a -> Exp a) ->
   Sym.Array p sh1 a ->
   (Exp ix0 -> Exp ix1) ->
   Sym.Array p sh0 a ->
   IO (p -> IO (Phys.Array sh1 a))
permute :: forall sh0 ix0 sh1 ix1 a p.
(C sh0, Index sh0 ~ ix0, C sh1, Index sh1 ~ ix1, C sh1, C a,
 C a) =>
(Exp a -> Exp a -> Exp a)
-> Array p sh1 a
-> (Exp ix0 -> Exp ix1)
-> Array p sh0 a
-> IO (p -> IO (Array sh1 a))
permute Exp a -> Exp a -> Exp a
accum Array p sh1 a
deflt Exp ix0 -> Exp ix1
ixmap Array p sh0 a
input =
   (Exp a -> Exp a -> Exp a)
-> Array p sh1 a
-> Array p sh0 (ix1, a)
-> IO (p -> IO (Array sh1 a))
forall sh0 ix0 sh1 ix1 a p.
(C sh0, Index sh0 ~ ix0, C sh1, Index sh1 ~ ix1, C sh1, C a,
 C a) =>
(Exp a -> Exp a -> Exp a)
-> Array p sh1 a
-> Array p sh0 (ix1, a)
-> IO (p -> IO (Array sh1 a))
scatter Exp a -> Exp a -> Exp a
accum Array p sh1 a
deflt
      ((Exp ix0 -> Exp a -> Exp (ix1, a))
-> Array p sh0 a -> Array p sh0 (ix1, a)
forall (array :: * -> * -> *) sh ix a b.
(C array, C sh, Index sh ~ ix) =>
(Exp ix -> Exp a -> Exp b) -> array sh a -> array sh b
Core.mapWithIndex ((T ix1 -> T a -> T (ix1, a)) -> Exp ix1 -> Exp a -> Exp (ix1, a)
forall a b c. (T a -> T b -> T c) -> Exp a -> Exp b -> Exp c
forall (val :: * -> *) a b c.
Value val =>
(T a -> T b -> T c) -> val a -> val b -> val c
Expr.lift2 T ix1 -> T a -> T (ix1, a)
forall a b. T a -> T b -> T (a, b)
MultiValue.zip (Exp ix1 -> Exp a -> Exp (ix1, a))
-> (Exp ix0 -> Exp ix1) -> Exp ix0 -> Exp a -> Exp (ix1, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp ix0 -> Exp ix1
ixmap) Array p sh0 a
input)