{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ForeignFunctionInterface #-}
module Data.Array.Knead.Symbolic.Physical (
   Array,
   shape,
   toList,
   fromList,
   vectorFromList,
   with,
   render,
   scanl1,
   mapAccumLSimple,
   scatter,
   scatterMaybe,
   permute,
   ) where

import qualified Data.Array.Knead.Symbolic.PhysicalPrivate as Priv
import qualified Data.Array.Knead.Symbolic.Private as Sym
import qualified Data.Array.Knead.Shape as Shape
import qualified Data.Array.Knead.Expression as Expr
import Data.Array.Knead.Symbolic.PhysicalPrivate (MarshalPtr)
import Data.Array.Knead.Code (getElementPtr)

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

import qualified Data.Array.Comfort.Storable.Mutable.Unchecked as MutArray
import qualified Data.Array.Comfort.Storable.Unchecked as Array
import qualified Data.Array.Comfort.Shape as ComfortShape
import Data.Array.Comfort.Storable.Unchecked (Array(Array))

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.Extra.Maybe as Maybe

import qualified LLVM.Core as LLVM

import Foreign.Storable (Storable, )
import Foreign.ForeignPtr (withForeignPtr, mallocForeignPtrArray, )
import Foreign.Ptr (FunPtr, Ptr, )

import Control.Monad.HT (void, (<=<), )
import Control.Applicative (liftA2, (<$>), )

import Prelude2010 hiding (scanl1)
import Prelude ()


shape :: Array sh a -> sh
shape :: forall sh a. Array sh a -> sh
shape = Array sh a -> sh
forall sh a. Array sh a -> sh
Array.shape

toList ::
   (Shape.C sh, Storable a) =>
   Array sh a -> IO [a]
toList :: forall sh a. (C sh, Storable a) => Array sh a -> IO [a]
toList = Array IO sh a -> IO [a]
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
Array m sh a -> m [a]
MutArray.toList (Array IO sh a -> IO [a])
-> (Array sh a -> IO (Array IO sh a)) -> Array sh a -> IO [a]
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Array sh a -> IO (Array IO sh a)
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
Array sh a -> m (Array m sh a)
MutArray.unsafeThaw

fromList ::
   (Shape.C sh, Storable a) =>
   sh -> [a] -> IO (Array sh a)
fromList :: forall sh a. (C sh, Storable a) => sh -> [a] -> IO (Array sh a)
fromList sh
sh = Array IO sh a -> IO (Array sh a)
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
Array m sh a -> m (Array sh a)
MutArray.unsafeFreeze (Array IO sh a -> IO (Array sh a))
-> ([a] -> IO (Array IO sh a)) -> [a] -> IO (Array sh a)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< sh -> [a] -> IO (Array IO sh a)
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
sh -> [a] -> m (Array m sh a)
MutArray.fromList sh
sh

vectorFromList ::
   (Num n, Storable a) =>
   [a] -> IO (Array (ComfortShape.ZeroBased n) a)
vectorFromList :: forall n a.
(Num n, Storable a) =>
[a] -> IO (Array (ZeroBased n) a)
vectorFromList [a]
xs =
   (ZeroBased Int -> ZeroBased n)
-> Array (ZeroBased Int) a -> Array (ZeroBased n) a
forall sh0 sh1 a. (sh0 -> sh1) -> Array sh0 a -> Array sh1 a
Array.mapShape (\(Shape.ZeroBased Int
n) -> n -> ZeroBased n
forall n. n -> ZeroBased n
Shape.ZeroBased (n -> ZeroBased n) -> n -> ZeroBased n
forall a b. (a -> b) -> a -> b
$ Int -> n
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) (Array (ZeroBased Int) a -> Array (ZeroBased n) a)
-> IO (Array (ZeroBased Int) a) -> IO (Array (ZeroBased n) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
   (Array IO (ZeroBased Int) a -> IO (Array (ZeroBased Int) a)
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
Array m sh a -> m (Array sh a)
MutArray.unsafeFreeze (Array IO (ZeroBased Int) a -> IO (Array (ZeroBased Int) a))
-> IO (Array IO (ZeroBased Int) a) -> IO (Array (ZeroBased Int) a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [a] -> IO (Array IO (ZeroBased Int) a)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
[a] -> m (Array m (ZeroBased Int) a)
MutArray.vectorFromList [a]
xs)


{- |
The symbolic array is only valid inside the enclosed action.
-}
with ::
   (Shape.C sh, Storable.C a) =>
   (Sym.Array sh a -> IO b) ->
   Array sh a -> IO b
with :: forall sh a b.
(C sh, C a) =>
(Array sh a -> IO b) -> Array sh a -> IO b
with Array sh a -> IO b
f (Array sh
sh ForeignPtr a
fptr) =
   ForeignPtr a -> (Ptr a -> IO b) -> IO b
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
fptr ((Ptr a -> IO b) -> IO b) -> (Ptr a -> IO b) -> IO b
forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr ->
      Array sh a -> IO b
f (Array sh a -> IO b) -> Array sh a -> IO b
forall a b. (a -> b) -> a -> b
$
      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
Sym.Array
         (sh -> Exp sh
forall sh (val :: * -> *). (C sh, Value val) => sh -> val sh
Shape.value sh
sh)
         (\Val (Index sh)
ix ->
            Value (Ptr a) -> Code r a
forall a r. C a => Value (Ptr a) -> CodeGenFunction r (T a)
forall r. Value (Ptr a) -> CodeGenFunction r (T a)
Storable.load (Value (Ptr a) -> Code r a)
-> CodeGenFunction r (Value (Ptr a)) -> Code r a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<
               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 (sh -> T sh
forall sh (val :: * -> *). (C sh, Value val) => sh -> val sh
Shape.value sh
sh) (Ptr a -> Value (Ptr a)
forall a. IsConst a => a -> Value a
LLVM.valueOf Ptr a
ptr) Val (Index sh)
ix)


type Importer f = FunPtr f -> f

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

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


materialize ::
   (Shape.C sh, Marshal.C sh, Storable.C a) =>
   String ->
   Exp sh ->
   (LLVM.Value (MarshalPtr sh) ->
    LLVM.Value (Ptr a) -> LLVM.CodeGenFunction () ()) ->
   IO (Array sh a)
materialize :: forall sh a.
(C sh, C sh, C a) =>
String
-> Exp sh
-> (Value (MarshalPtr sh)
    -> Value (Ptr a) -> CodeGenFunction () ())
-> IO (Array sh a)
materialize String
name Exp sh
esh Value (Ptr (Struct (Repr sh)))
-> Value (Ptr a) -> CodeGenFunction () ()
code =
   (Ptr (Struct (Repr sh)) -> IO (Array sh a)) -> IO (Array sh a)
forall a b. IsType a => (Ptr a -> IO b) -> IO b
Marshal.alloca ((Ptr (Struct (Repr sh)) -> IO (Array sh a)) -> IO (Array sh a))
-> (Ptr (Struct (Repr sh)) -> IO (Array sh a)) -> IO (Array sh a)
forall a b. (a -> b) -> a -> b
$ \Ptr (Struct (Repr sh))
lshptr -> do
      (Ptr (Struct (Repr sh)) -> IO Size
fsh, Ptr (Struct (Repr sh)) -> Ptr a -> IO ()
farr) <-
         String
-> Exec
     (Ptr (Struct (Repr sh)) -> IO Size,
      Ptr (Struct (Repr sh)) -> Ptr a -> IO ())
-> IO
     (Ptr (Struct (Repr sh)) -> IO Size,
      Ptr (Struct (Repr sh)) -> Ptr a -> IO ())
forall funcs. String -> Exec funcs -> IO funcs
Code.compile String
name (Exec
   (Ptr (Struct (Repr sh)) -> IO Size,
    Ptr (Struct (Repr sh)) -> Ptr a -> IO ())
 -> IO
      (Ptr (Struct (Repr sh)) -> IO Size,
       Ptr (Struct (Repr sh)) -> Ptr a -> IO ()))
-> Exec
     (Ptr (Struct (Repr sh)) -> IO Size,
      Ptr (Struct (Repr sh)) -> Ptr a -> IO ())
-> IO
     (Ptr (Struct (Repr sh)) -> IO Size,
      Ptr (Struct (Repr sh)) -> Ptr a -> IO ())
forall a b. (a -> b) -> a -> b
$
         ((Ptr (Struct (Repr sh)) -> IO Size)
 -> (Ptr (Struct (Repr sh)) -> Ptr a -> IO ())
 -> (Ptr (Struct (Repr sh)) -> IO Size,
     Ptr (Struct (Repr sh)) -> Ptr a -> IO ()))
-> Compose
     CodeGenModule EngineAccess (Ptr (Struct (Repr sh)) -> IO Size)
-> Compose
     CodeGenModule
     EngineAccess
     (Ptr (Struct (Repr sh)) -> Ptr a -> IO ())
-> Exec
     (Ptr (Struct (Repr sh)) -> IO Size,
      Ptr (Struct (Repr sh)) -> Ptr a -> IO ())
forall a b c.
(a -> b -> c)
-> Compose CodeGenModule EngineAccess a
-> Compose CodeGenModule EngineAccess b
-> Compose CodeGenModule EngineAccess c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (,)
            (Importer (Ptr (Struct (Repr sh)) -> IO Size)
-> String
-> CodeGen (Ptr (Struct (Repr sh)) -> IO Size)
-> Compose
     CodeGenModule EngineAccess (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 sh)) -> IO Size)
forall sh. Importer (Ptr sh -> IO Size)
callShaper String
"shape" (CodeGen (Ptr (Struct (Repr sh)) -> IO Size)
 -> Compose
      CodeGenModule EngineAccess (Ptr (Struct (Repr sh)) -> IO Size))
-> CodeGen (Ptr (Struct (Repr sh)) -> IO Size)
-> Compose
     CodeGenModule EngineAccess (Ptr (Struct (Repr sh)) -> IO Size)
forall a b. (a -> b) -> a -> b
$ \Value (Ptr (Struct (Repr sh)))
ptr -> 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)))
ptr
               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)
            (Importer (Ptr (Struct (Repr sh)) -> Ptr a -> IO ())
-> String
-> CodeGen (Ptr (Struct (Repr sh)) -> Ptr a -> IO ())
-> Compose
     CodeGenModule
     EngineAccess
     (Ptr (Struct (Repr sh)) -> Ptr a -> IO ())
forall f.
(ExecutionFunction f, C f) =>
Importer f -> String -> CodeGen f -> Exec f
Code.createFunction Importer (Ptr (Struct (Repr sh)) -> Ptr a -> IO ())
forall sh a. Importer (Ptr sh -> Ptr a -> IO ())
callRenderer String
"fill" CodeGen (Ptr (Struct (Repr sh)) -> Ptr a -> IO ())
Value (Ptr (Struct (Repr sh)))
-> Value (Ptr a) -> CodeGenFunction () ()
code)
      Size
n <- Ptr (Struct (Repr sh)) -> IO Size
fsh Ptr (Struct (Repr sh))
lshptr
      ForeignPtr a
fptr <- Int -> IO (ForeignPtr a)
forall a. Storable a => Int -> IO (ForeignPtr a)
mallocForeignPtrArray (Size -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Size
n)
      ForeignPtr a -> (Ptr a -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
fptr ((Ptr a -> IO ()) -> IO ()) -> (Ptr a -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr (Struct (Repr sh)) -> Ptr a -> IO ()
farr Ptr (Struct (Repr sh))
lshptr
      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))
lshptr
      Array sh a -> IO (Array sh a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (sh -> ForeignPtr a -> Array sh a
forall sh a. sh -> ForeignPtr a -> Array sh a
Array sh
sh ForeignPtr a
fptr)

render ::
   (Shape.C sh, Marshal.C sh, Storable.C a) =>
   Sym.Array sh a -> IO (Array sh a)
render :: forall sh a. (C sh, C sh, C a) => Array sh a -> IO (Array sh a)
render (Sym.Array Exp sh
esh forall r. Val (Index sh) -> Code r a
code) =
   String
-> Exp sh
-> (Value (Ptr (Struct (Repr sh)))
    -> Value (Ptr a) -> CodeGenFunction () ())
-> IO (Array sh a)
forall sh a.
(C sh, C sh, C a) =>
String
-> Exp sh
-> (Value (MarshalPtr sh)
    -> Value (Ptr a) -> CodeGenFunction () ())
-> IO (Array sh a)
materialize String
"render" Exp sh
esh ((Value (Ptr (Struct (Repr sh)))
  -> Value (Ptr a) -> CodeGenFunction () ())
 -> IO (Array sh a))
-> (Value (Ptr (Struct (Repr sh)))
    -> Value (Ptr a) -> CodeGenFunction () ())
-> IO (Array sh a)
forall a b. (a -> b) -> a -> b
$ \Value (Ptr (Struct (Repr sh)))
sptr Value (Ptr a)
ptr -> do
      let step :: Val (Index sh)
-> Value (Ptr a) -> CodeGenFunction () (Value (Ptr a))
step Val (Index sh)
ix Value (Ptr a)
p = (T a -> Value (Ptr a) -> CodeGenFunction () (Value (Ptr a)))
-> Value (Ptr a) -> T a -> CodeGenFunction () (Value (Ptr a))
forall a b c. (a -> b -> c) -> b -> a -> c
flip T a -> Value (Ptr a) -> CodeGenFunction () (Value (Ptr a))
forall a ptr r.
(C a, Value (Ptr a) ~ ptr) =>
T a -> ptr -> CodeGenFunction r ptr
Storable.storeNext Value (Ptr a)
p (T a -> CodeGenFunction () (Value (Ptr a)))
-> CodeGenFunction () (T a) -> CodeGenFunction () (Value (Ptr a))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Val (Index sh) -> CodeGenFunction () (T a)
forall r. Val (Index sh) -> Code r a
code Val (Index sh)
ix
      T sh
sh <- Exp sh
-> Value (Ptr (Struct (Repr sh))) -> CodeGenFunction () (T sh)
forall sh (f :: * -> *) r.
C sh =>
f sh -> Value (Ptr (Struct sh)) -> CodeGenFunction r (T sh)
Shape.load Exp sh
esh Value (Ptr (Struct (Repr sh)))
sptr
      CodeGenFunction () (Value (Ptr a)) -> CodeGenFunction () ()
forall (m :: * -> *) a. Monad m => m a -> m ()
void (CodeGenFunction () (Value (Ptr a)) -> CodeGenFunction () ())
-> CodeGenFunction () (Value (Ptr a)) -> CodeGenFunction () ()
forall a b. (a -> b) -> a -> b
$ (Val (Index sh)
 -> Value (Ptr a) -> CodeGenFunction () (Value (Ptr a)))
-> T sh -> Value (Ptr a) -> CodeGenFunction () (Value (Ptr a))
forall sh ix state r.
(C sh, Index sh ~ ix, Phi state) =>
(T ix -> state -> CodeGenFunction r state)
-> T sh -> state -> CodeGenFunction r state
forall ix state r.
(Index sh ~ ix, Phi state) =>
(T ix -> state -> CodeGenFunction r state)
-> T sh -> state -> CodeGenFunction r state
Shape.loop Val (Index sh)
-> Value (Ptr a) -> CodeGenFunction () (Value (Ptr a))
step T sh
sh Value (Ptr a)
ptr

scanl1 ::
   (Shape.C sh, Marshal.C sh,
    Shape.C n, Marshal.C n,
    Storable.C a, MultiValue.C a) =>
   (Exp a -> Exp a -> Exp a) ->
   Sym.Array (sh, n) a -> IO (Array (sh, n) a)
scanl1 :: forall sh n a.
(C sh, C sh, C n, C n, C a, C a) =>
(Exp a -> Exp a -> Exp a)
-> Array (sh, n) a -> IO (Array (sh, n) a)
scanl1 Exp a -> Exp a -> Exp a
f (Sym.Array Exp (sh, n)
esh forall r. Val (Index (sh, n)) -> Code r a
code) =
   String
-> Exp (sh, n)
-> (Value (Ptr (Struct (Repr (sh, n))))
    -> Value (Ptr a) -> CodeGenFunction () ())
-> IO (Array (sh, n) a)
forall sh a.
(C sh, C sh, C a) =>
String
-> Exp sh
-> (Value (MarshalPtr sh)
    -> Value (Ptr a) -> CodeGenFunction () ())
-> IO (Array sh a)
materialize String
"scanl1" Exp (sh, n)
esh ((Value (Ptr (Struct (Repr (sh, n))))
  -> Value (Ptr a) -> CodeGenFunction () ())
 -> IO (Array (sh, n) a))
-> (Value (Ptr (Struct (Repr (sh, n))))
    -> Value (Ptr a) -> CodeGenFunction () ())
-> IO (Array (sh, n) a)
forall a b. (a -> b) -> a -> b
$ \Value (Ptr (Struct (Repr (sh, n))))
sptr Value (Ptr a)
ptr -> do
      (T sh
sh, T n
n) <- T (sh, n) -> (T sh, T n)
forall a b. T (a, b) -> (T a, T b)
MultiValue.unzip (T (sh, n) -> (T sh, T n))
-> CodeGenFunction () (T (sh, n)) -> CodeGenFunction () (T sh, T n)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp (sh, n)
-> Value (Ptr (Struct (Repr (sh, n))))
-> CodeGenFunction () (T (sh, n))
forall sh (f :: * -> *) r.
C sh =>
f sh -> Value (Ptr (Struct sh)) -> CodeGenFunction r (T sh)
Shape.load Exp (sh, n)
esh Value (Ptr (Struct (Repr (sh, n))))
sptr
      let step :: T (Index sh) -> Value (Ptr a) -> CodeGenFunction () (Value (Ptr a))
step T (Index sh)
ix Value (Ptr a)
ptrStart =
             ((Value (Ptr a), T (Val a)) -> Value (Ptr a))
-> CodeGenFunction () (Value (Ptr a), T (Val a))
-> CodeGenFunction () (Value (Ptr a))
forall a b.
(a -> b) -> CodeGenFunction () a -> CodeGenFunction () b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Value (Ptr a), T (Val a)) -> Value (Ptr a)
forall a b. (a, b) -> a
fst (CodeGenFunction () (Value (Ptr a), T (Val a))
 -> CodeGenFunction () (Value (Ptr a)))
-> CodeGenFunction () (Value (Ptr a), T (Val a))
-> CodeGenFunction () (Value (Ptr a))
forall a b. (a -> b) -> a -> b
$
             (\T (Index n)
-> (Value (Ptr a), T (Val a))
-> CodeGenFunction () (Value (Ptr a), T (Val a))
body -> (T (Index n)
 -> (Value (Ptr a), T (Val a))
 -> CodeGenFunction () (Value (Ptr a), T (Val a)))
-> T n
-> (Value (Ptr a), T (Val a))
-> CodeGenFunction () (Value (Ptr a), T (Val a))
forall sh ix state r.
(C sh, Index sh ~ ix, Phi state) =>
(T ix -> state -> CodeGenFunction r state)
-> T sh -> state -> CodeGenFunction r state
forall ix state r.
(Index n ~ ix, Phi state) =>
(T ix -> state -> CodeGenFunction r state)
-> T n -> state -> CodeGenFunction r state
Shape.loop T (Index n)
-> (Value (Ptr a), T (Val a))
-> CodeGenFunction () (Value (Ptr a), T (Val a))
body T n
n (Value (Ptr a)
ptrStart, T (Val a)
forall a. Undefined a => T a
Maybe.nothing)) ((T (Index n)
  -> (Value (Ptr a), T (Val a))
  -> CodeGenFunction () (Value (Ptr a), T (Val a)))
 -> CodeGenFunction () (Value (Ptr a), T (Val a)))
-> (T (Index n)
    -> (Value (Ptr a), T (Val a))
    -> CodeGenFunction () (Value (Ptr a), T (Val a)))
-> CodeGenFunction () (Value (Ptr a), T (Val a))
forall a b. (a -> b) -> a -> b
$
                   \T (Index n)
k0 (Value (Ptr a)
ptr0, T (Val a)
macc0) -> do
                Val a
a <- Val (Index (sh, n)) -> Code () a
forall r. Val (Index (sh, n)) -> Code r a
code (Val (Index (sh, n)) -> Code () a)
-> Val (Index (sh, n)) -> Code () a
forall a b. (a -> b) -> a -> b
$ T (Index sh) -> T (Index n) -> T (Index sh, Index n)
forall a b. T a -> T b -> T (a, b)
MultiValue.zip T (Index sh)
ix T (Index n)
k0
                Val a
acc1 <- T (Val a) -> Code () a -> (Val a -> Code () a) -> Code () a
forall b a r.
Phi b =>
T a
-> CodeGenFunction r b
-> (a -> CodeGenFunction r b)
-> CodeGenFunction r b
Maybe.run T (Val a)
macc0 (Val a -> Code () a
forall a. a -> CodeGenFunction () a
forall (m :: * -> *) a. Monad m => a -> m a
return Val a
a) ((Val a -> Val a -> Code () a) -> Val a -> Val a -> Code () a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Exp a -> Exp a -> Exp a) -> Val a -> Val a -> Code () a
forall ae am be bm ce cm r.
(Aggregate ae am, Aggregate be bm, Aggregate ce cm) =>
(ae -> be -> ce) -> am -> bm -> CodeGenFunction r cm
Expr.unliftM2 Exp a -> Exp a -> Exp a
f) Val a
a)
                Value (Ptr a)
ptr1 <- Val a -> Value (Ptr a) -> CodeGenFunction () (Value (Ptr a))
forall a ptr r.
(C a, Value (Ptr a) ~ ptr) =>
T a -> ptr -> CodeGenFunction r ptr
Storable.storeNext Val a
acc1 Value (Ptr a)
ptr0
                (Value (Ptr a), T (Val a))
-> CodeGenFunction () (Value (Ptr a), T (Val a))
forall a. a -> CodeGenFunction () a
forall (m :: * -> *) a. Monad m => a -> m a
return (Value (Ptr a)
ptr1, Val a -> T (Val a)
forall a. a -> T a
Maybe.just Val a
acc1)
      CodeGenFunction () (Value (Ptr a)) -> CodeGenFunction () ()
forall (m :: * -> *) a. Monad m => m a -> m ()
void (CodeGenFunction () (Value (Ptr a)) -> CodeGenFunction () ())
-> CodeGenFunction () (Value (Ptr a)) -> CodeGenFunction () ()
forall a b. (a -> b) -> a -> b
$ (T (Index sh)
 -> Value (Ptr a) -> CodeGenFunction () (Value (Ptr a)))
-> T sh -> Value (Ptr a) -> CodeGenFunction () (Value (Ptr a))
forall sh ix state r.
(C sh, Index sh ~ ix, Phi state) =>
(T ix -> state -> CodeGenFunction r state)
-> T sh -> state -> CodeGenFunction r state
forall ix state r.
(Index sh ~ ix, Phi state) =>
(T ix -> state -> CodeGenFunction r state)
-> T sh -> state -> CodeGenFunction r state
Shape.loop T (Index sh) -> Value (Ptr a) -> CodeGenFunction () (Value (Ptr a))
step T sh
sh Value (Ptr a)
ptr

mapAccumLSimple ::
   (Shape.C sh, Marshal.C sh,
    Shape.C n, Marshal.C n,
    MultiValue.C acc, Storable.C x, Storable.C y) =>
   (Exp acc -> Exp x -> Exp (acc,y)) ->
   Sym.Array sh acc -> Sym.Array (sh, n) x -> IO (Array (sh, n) y)
mapAccumLSimple :: forall sh n acc x y.
(C sh, C sh, C n, C n, C acc, C x, C y) =>
(Exp acc -> Exp x -> Exp (acc, y))
-> Array sh acc -> Array (sh, n) x -> IO (Array (sh, n) y)
mapAccumLSimple Exp acc -> Exp x -> Exp (acc, y)
f Array sh acc
arrInit Array (sh, n) x
arrData =
   String
-> Exp (sh, n)
-> (Value (Ptr (Struct (Repr (sh, n))))
    -> Value (Ptr y) -> CodeGenFunction () ())
-> IO (Array (sh, n) y)
forall sh a.
(C sh, C sh, C a) =>
String
-> Exp sh
-> (Value (MarshalPtr sh)
    -> Value (Ptr a) -> CodeGenFunction () ())
-> IO (Array sh a)
materialize String
"mapAccumLSimple" (Array (sh, n) x -> Exp (sh, n)
forall sh a. Array sh a -> Exp sh
Sym.shape Array (sh, n) x
arrData) ((Value (Ptr (Struct (Repr (sh, n))))
  -> Value (Ptr y) -> CodeGenFunction () ())
 -> IO (Array (sh, n) y))
-> (Value (Ptr (Struct (Repr (sh, n))))
    -> Value (Ptr y) -> CodeGenFunction () ())
-> IO (Array (sh, n) y)
forall a b. (a -> b) -> a -> b
$
      (Exp acc -> Exp x -> Exp (acc, y))
-> Array sh acc
-> Array (sh, n) x
-> Value (Ptr (Struct (Repr (sh, n))))
-> Value (Ptr y)
-> CodeGenFunction () ()
forall sh n acc x y r.
(C sh, C sh, C n, C n, C acc, C x, C y) =>
(Exp acc -> Exp x -> Exp (acc, y))
-> Array sh acc
-> Array (sh, n) x
-> Value (MarshalPtr (sh, n))
-> Value (Ptr y)
-> CodeGenFunction r ()
Priv.mapAccumLSimple Exp acc -> Exp x -> Exp (acc, y)
f Array sh acc
arrInit Array (sh, n) x
arrData

scatterMaybe ::
   (Shape.C sh0, Shape.Index sh0 ~ ix0,
    Shape.C sh1, Shape.Index sh1 ~ ix1, Marshal.C sh1,
    Storable.C a) =>
   (Exp a -> Exp a -> Exp a) ->
   Sym.Array sh1 a ->
   Sym.Array sh0 (Maybe (ix1, a)) -> IO (Array sh1 a)
scatterMaybe :: forall sh0 ix0 sh1 ix1 a.
(C sh0, Index sh0 ~ ix0, C sh1, Index sh1 ~ ix1, C sh1, C a) =>
(Exp a -> Exp a -> Exp a)
-> Array sh1 a -> Array sh0 (Maybe (ix1, a)) -> IO (Array sh1 a)
scatterMaybe Exp a -> Exp a -> Exp a
accum Array sh1 a
arrInit Array sh0 (Maybe (ix1, a))
arrMap =
   String
-> Exp sh1
-> (Value (Ptr (Struct (Repr sh1)))
    -> Value (Ptr a) -> CodeGenFunction () ())
-> IO (Array sh1 a)
forall sh a.
(C sh, C sh, C a) =>
String
-> Exp sh
-> (Value (MarshalPtr sh)
    -> Value (Ptr a) -> CodeGenFunction () ())
-> IO (Array sh a)
materialize String
"scatterMaybe" (Array sh1 a -> Exp sh1
forall sh a. Array sh a -> Exp sh
Sym.shape Array sh1 a
arrInit) ((Value (Ptr (Struct (Repr sh1)))
  -> Value (Ptr a) -> CodeGenFunction () ())
 -> IO (Array sh1 a))
-> (Value (Ptr (Struct (Repr sh1)))
    -> Value (Ptr a) -> CodeGenFunction () ())
-> IO (Array sh1 a)
forall a b. (a -> b) -> a -> b
$
      (Exp a -> Exp a -> Exp a)
-> Array sh1 a
-> Array sh0 (Maybe (ix1, a))
-> Value (Ptr (Struct (Repr sh1)))
-> Value (Ptr a)
-> CodeGenFunction () ()
forall sh0 ix0 sh1 ix1 a r.
(C sh0, Index sh0 ~ ix0, C sh1, Index sh1 ~ ix1, C sh1, C a) =>
(Exp a -> Exp a -> Exp a)
-> Array sh1 a
-> Array sh0 (Maybe (ix1, a))
-> Value (MarshalPtr sh1)
-> Value (Ptr a)
-> CodeGenFunction r ()
Priv.scatterMaybe Exp a -> Exp a -> Exp a
accum Array sh1 a
arrInit Array sh0 (Maybe (ix1, a))
arrMap

scatter ::
   (Shape.C sh0, Shape.Index sh0 ~ ix0,
    Shape.C sh1, Shape.Index sh1 ~ ix1, Marshal.C sh1,
    Storable.C a) =>
   (Exp a -> Exp a -> Exp a) ->
   Sym.Array sh1 a ->
   Sym.Array sh0 (ix1, a) -> IO (Array sh1 a)
scatter :: forall sh0 ix0 sh1 ix1 a.
(C sh0, Index sh0 ~ ix0, C sh1, Index sh1 ~ ix1, C sh1, C a) =>
(Exp a -> Exp a -> Exp a)
-> Array sh1 a -> Array sh0 (ix1, a) -> IO (Array sh1 a)
scatter Exp a -> Exp a -> Exp a
accum Array sh1 a
arrInit Array sh0 (ix1, a)
arrMap =
   String
-> Exp sh1
-> (Value (Ptr (Struct (Repr sh1)))
    -> Value (Ptr a) -> CodeGenFunction () ())
-> IO (Array sh1 a)
forall sh a.
(C sh, C sh, C a) =>
String
-> Exp sh
-> (Value (MarshalPtr sh)
    -> Value (Ptr a) -> CodeGenFunction () ())
-> IO (Array sh a)
materialize String
"scatter" (Array sh1 a -> Exp sh1
forall sh a. Array sh a -> Exp sh
Sym.shape Array sh1 a
arrInit) ((Value (Ptr (Struct (Repr sh1)))
  -> Value (Ptr a) -> CodeGenFunction () ())
 -> IO (Array sh1 a))
-> (Value (Ptr (Struct (Repr sh1)))
    -> Value (Ptr a) -> CodeGenFunction () ())
-> IO (Array sh1 a)
forall a b. (a -> b) -> a -> b
$
      (Exp a -> Exp a -> Exp a)
-> Array sh1 a
-> Array sh0 (Index sh1, a)
-> Value (Ptr (Struct (Repr sh1)))
-> Value (Ptr a)
-> CodeGenFunction () ()
forall sh0 ix0 sh1 ix1 a r.
(C sh0, Index sh0 ~ ix0, C sh1, Index sh1 ~ ix1, C sh1, C a) =>
(Exp a -> Exp a -> Exp a)
-> Array sh1 a
-> Array sh0 (Index sh1, a)
-> Value (MarshalPtr sh1)
-> Value (Ptr a)
-> CodeGenFunction r ()
Priv.scatter Exp a -> Exp a -> Exp a
accum Array sh1 a
arrInit Array sh0 (ix1, a)
Array sh0 (Index sh1, a)
arrMap

permute ::
   (Shape.C sh0, Shape.Index sh0 ~ ix0,
    Shape.C sh1, Shape.Index sh1 ~ ix1, Marshal.C sh1,
    Storable.C a) =>
   (Exp a -> Exp a -> Exp a) ->
   Sym.Array sh1 a ->
   (Exp ix0 -> Exp ix1) ->
   Sym.Array sh0 a ->
   IO (Array sh1 a)
permute :: forall sh0 ix0 sh1 ix1 a.
(C sh0, Index sh0 ~ ix0, C sh1, Index sh1 ~ ix1, C sh1, C a) =>
(Exp a -> Exp a -> Exp a)
-> Array sh1 a
-> (Exp ix0 -> Exp ix1)
-> Array sh0 a
-> IO (Array sh1 a)
permute Exp a -> Exp a -> Exp a
accum Array sh1 a
deflt Exp ix0 -> Exp ix1
ixmap Array sh0 a
input =
   (Exp a -> Exp a -> Exp a)
-> Array sh1 a -> Array sh0 (ix1, a) -> IO (Array sh1 a)
forall sh0 ix0 sh1 ix1 a.
(C sh0, Index sh0 ~ ix0, C sh1, Index sh1 ~ ix1, C sh1, C a) =>
(Exp a -> Exp a -> Exp a)
-> Array sh1 a -> Array sh0 (ix1, a) -> IO (Array sh1 a)
scatter Exp a -> Exp a -> Exp a
accum Array sh1 a
deflt
      ((Exp ix0 -> Exp a -> Exp (ix1, a))
-> Array sh0 a -> Array 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
Sym.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 sh0 a
input)