{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Data.Array.Knead.Symbolic.PhysicalPrivate where

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.Code (getElementPtr)

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.Control as C

import qualified LLVM.Core as LLVM

import Foreign.Ptr (Ptr, )

import qualified Control.Applicative.HT as App
import Control.Monad.HT (void, )
import Control.Applicative ((<$>), )

import Data.Tuple.HT (mapSnd, )

import Prelude2010
import Prelude ()



type MarshalPtr a = LLVM.Ptr (Marshal.Struct a)

writeArray ::
   (Shape.C sh, Shape.Index sh ~ ix, Storable.C a) =>
   MultiValue.T sh ->
   (MultiValue.T ix -> LLVM.CodeGenFunction r (MultiValue.T a)) ->
   LLVM.Value (Ptr a) ->
   LLVM.CodeGenFunction r (LLVM.Value (Ptr a))
writeArray :: forall sh ix a r.
(C sh, Index sh ~ ix, C a) =>
T sh
-> (T ix -> CodeGenFunction r (T a))
-> Value (Ptr a)
-> CodeGenFunction r (Value (Ptr a))
writeArray T sh
sh T ix -> CodeGenFunction r (T a)
code Value (Ptr a)
ptr = do
   let clear :: T ix -> Value (Ptr a) -> CodeGenFunction r (Value (Ptr a))
clear T ix
ix Value (Ptr a)
p = (T a -> Value (Ptr a) -> CodeGenFunction r (Value (Ptr a)))
-> Value (Ptr a) -> T a -> CodeGenFunction r (Value (Ptr a))
forall a b c. (a -> b -> c) -> b -> a -> c
flip T a -> Value (Ptr a) -> CodeGenFunction r (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 r (Value (Ptr a)))
-> CodeGenFunction r (T a) -> CodeGenFunction r (Value (Ptr a))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< T ix -> CodeGenFunction r (T a)
code T ix
ix
   (T ix -> Value (Ptr a) -> CodeGenFunction r (Value (Ptr a)))
-> T sh -> Value (Ptr a) -> CodeGenFunction r (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 ix -> Value (Ptr a) -> CodeGenFunction r (Value (Ptr a))
clear T sh
sh Value (Ptr a)
ptr


mapAccumLLoop ::
   (MultiValue.C acc, Storable.C b,
    Shape.C sh, Shape.Index sh ~ ix) =>
   (MultiValue.T ix -> LLVM.CodeGenFunction r (MultiValue.T a)) ->
   (Exp acc -> Exp a -> Exp (acc, b)) ->
   MultiValue.T sh ->
   LLVM.Value (Ptr b) -> MultiValue.T acc ->
   LLVM.CodeGenFunction r (LLVM.Value (Ptr b), MultiValue.T acc)
mapAccumLLoop :: forall acc b sh ix r a.
(C acc, C b, C sh, Index sh ~ ix) =>
(T ix -> CodeGenFunction r (T a))
-> (Exp acc -> Exp a -> Exp (acc, b))
-> T sh
-> Value (Ptr b)
-> T acc
-> CodeGenFunction r (Value (Ptr b), T acc)
mapAccumLLoop T ix -> CodeGenFunction r (T a)
code Exp acc -> Exp a -> Exp (acc, b)
f T sh
n Value (Ptr b)
yPtr T acc
accInit = do
   let step :: T ix
-> (Value (Ptr b), T acc)
-> CodeGenFunction r (Value (Ptr b), T acc)
step T ix
k0 (Value (Ptr b)
ptr0, T acc
acc0) = do
         T a
x <- T ix -> CodeGenFunction r (T a)
code T ix
k0
         (T acc
acc1,T b
y) <- T (acc, b) -> (T acc, T b)
forall a b. T (a, b) -> (T a, T b)
MultiValue.unzip (T (acc, b) -> (T acc, T b))
-> CodeGenFunction r (T (acc, b)) -> CodeGenFunction r (T acc, T b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp acc -> Exp a -> Exp (acc, b))
-> T acc -> T a -> CodeGenFunction r (T (acc, b))
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 acc -> Exp a -> Exp (acc, b)
f T acc
acc0 T a
x
         Value (Ptr b)
ptr1 <- T b -> Value (Ptr b) -> CodeGenFunction r (Value (Ptr b))
forall a ptr r.
(C a, Value (Ptr a) ~ ptr) =>
T a -> ptr -> CodeGenFunction r ptr
Storable.storeNext T b
y Value (Ptr b)
ptr0
         (Value (Ptr b), T acc) -> CodeGenFunction r (Value (Ptr b), T acc)
forall a. a -> CodeGenFunction r a
forall (m :: * -> *) a. Monad m => a -> m a
return (Value (Ptr b)
ptr1, T acc
acc1)
   (T ix
 -> (Value (Ptr b), T acc)
 -> CodeGenFunction r (Value (Ptr b), T acc))
-> T sh
-> (Value (Ptr b), T acc)
-> CodeGenFunction r (Value (Ptr b), T acc)
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 ix
-> (Value (Ptr b), T acc)
-> CodeGenFunction r (Value (Ptr b), T acc)
step T sh
n (Value (Ptr b)
yPtr, T acc
accInit)

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 ->
   LLVM.Value (MarshalPtr (sh,n)) ->
   LLVM.Value (Ptr y) ->
   LLVM.CodeGenFunction r ()
mapAccumLSimple :: 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 ()
mapAccumLSimple Exp acc -> Exp x -> Exp (acc, y)
f (Sym.Array Exp sh
_ forall r. Val (Index sh) -> Code r acc
initCode) (Sym.Array Exp (sh, n)
esh forall r. Val (Index (sh, n)) -> Code r x
code) Value (MarshalPtr (sh, n))
sptr Value (Ptr y)
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 r (T (sh, n)) -> CodeGenFunction r (T sh, T n)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp (sh, n)
-> Value (MarshalPtr (sh, n)) -> CodeGenFunction r (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 (MarshalPtr (sh, n))
sptr
   let step :: Val (Index sh)
-> Value (Ptr y) -> CodeGenFunction r (Value (Ptr y))
step Val (Index sh)
ix Value (Ptr y)
ptrStart = do
         Val acc
accInit <- Val (Index sh) -> Code r acc
forall r. Val (Index sh) -> Code r acc
initCode Val (Index sh)
ix
         (Value (Ptr y), Val acc) -> Value (Ptr y)
forall a b. (a, b) -> a
fst ((Value (Ptr y), Val acc) -> Value (Ptr y))
-> CodeGenFunction r (Value (Ptr y), Val acc)
-> CodeGenFunction r (Value (Ptr y))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (T (Index n) -> CodeGenFunction r (T x))
-> (Exp acc -> Exp x -> Exp (acc, y))
-> T n
-> Value (Ptr y)
-> Val acc
-> CodeGenFunction r (Value (Ptr y), Val acc)
forall acc b sh ix r a.
(C acc, C b, C sh, Index sh ~ ix) =>
(T ix -> CodeGenFunction r (T a))
-> (Exp acc -> Exp a -> Exp (acc, b))
-> T sh
-> Value (Ptr b)
-> T acc
-> CodeGenFunction r (Value (Ptr b), T acc)
mapAccumLLoop (T (Index sh, Index n) -> CodeGenFunction r (T x)
Val (Index (sh, n)) -> CodeGenFunction r (T x)
forall r. Val (Index (sh, n)) -> Code r x
code (T (Index sh, Index n) -> CodeGenFunction r (T x))
-> (T (Index n) -> T (Index sh, Index n))
-> T (Index n)
-> CodeGenFunction r (T x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Val (Index sh) -> T (Index n) -> T (Index sh, Index n)
forall a b. T a -> T b -> T (a, b)
MultiValue.zip Val (Index sh)
ix) Exp acc -> Exp x -> Exp (acc, y)
f T n
n Value (Ptr y)
ptrStart Val acc
accInit
   CodeGenFunction r (Value (Ptr y)) -> CodeGenFunction r ()
forall (m :: * -> *) a. Monad m => m a -> m ()
void (CodeGenFunction r (Value (Ptr y)) -> CodeGenFunction r ())
-> CodeGenFunction r (Value (Ptr y)) -> CodeGenFunction r ()
forall a b. (a -> b) -> a -> b
$ (Val (Index sh)
 -> Value (Ptr y) -> CodeGenFunction r (Value (Ptr y)))
-> T sh -> Value (Ptr y) -> CodeGenFunction r (Value (Ptr y))
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 y) -> CodeGenFunction r (Value (Ptr y))
step T sh
sh Value (Ptr y)
ptr

mapAccumLSequence ::
   (Shape.C n, Marshal.C n,
    MultiValue.C acc, Storable.C final,
    Storable.C x,
    Storable.C y) =>
   (Exp acc -> Exp x -> Exp (acc,y)) ->
   (Exp acc -> Exp final) ->
   Exp acc -> Sym.Array n x ->
   LLVM.Value (Ptr final) ->
   LLVM.Value (MarshalPtr n) ->
   LLVM.Value (Ptr y) ->
   LLVM.CodeGenFunction r ()
mapAccumLSequence :: forall n acc final x y r.
(C n, C n, C acc, C final, C x, C y) =>
(Exp acc -> Exp x -> Exp (acc, y))
-> (Exp acc -> Exp final)
-> Exp acc
-> Array n x
-> Value (Ptr final)
-> Value (MarshalPtr n)
-> Value (Ptr y)
-> CodeGenFunction r ()
mapAccumLSequence Exp acc -> Exp x -> Exp (acc, y)
f Exp acc -> Exp final
final Exp acc
initExp (Sym.Array Exp n
esh forall r. Val (Index n) -> Code r x
code) Value (Ptr final)
accPtr Value (MarshalPtr n)
sptr Value (Ptr y)
yPtr = do
   T n
n <- Exp n -> Value (MarshalPtr n) -> CodeGenFunction r (T n)
forall sh (f :: * -> *) r.
C sh =>
f sh -> Value (Ptr (Struct sh)) -> CodeGenFunction r (T sh)
Shape.load Exp n
esh Value (MarshalPtr n)
sptr
   T acc
accInit <- Exp acc -> forall r. CodeGenFunction r (T acc)
forall a. Exp a -> forall r. CodeGenFunction r (T a)
unExp Exp acc
initExp
   T acc
accExit <- (Value (Ptr y), T acc) -> T acc
forall a b. (a, b) -> b
snd ((Value (Ptr y), T acc) -> T acc)
-> CodeGenFunction r (Value (Ptr y), T acc)
-> CodeGenFunction r (T acc)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Val (Index n) -> CodeGenFunction r (T x))
-> (Exp acc -> Exp x -> Exp (acc, y))
-> T n
-> Value (Ptr y)
-> T acc
-> CodeGenFunction r (Value (Ptr y), T acc)
forall acc b sh ix r a.
(C acc, C b, C sh, Index sh ~ ix) =>
(T ix -> CodeGenFunction r (T a))
-> (Exp acc -> Exp a -> Exp (acc, b))
-> T sh
-> Value (Ptr b)
-> T acc
-> CodeGenFunction r (Value (Ptr b), T acc)
mapAccumLLoop Val (Index n) -> CodeGenFunction r (T x)
forall r. Val (Index n) -> Code r x
code Exp acc -> Exp x -> Exp (acc, y)
f T n
n Value (Ptr y)
yPtr T acc
accInit
   (T final -> Value (Ptr final) -> CodeGenFunction r ())
-> Value (Ptr final) -> T final -> CodeGenFunction r ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip T final -> Value (Ptr final) -> CodeGenFunction r ()
forall r. T final -> Value (Ptr final) -> CodeGenFunction r ()
forall a r. C a => T a -> Value (Ptr a) -> CodeGenFunction r ()
Storable.store Value (Ptr final)
accPtr (T final -> CodeGenFunction r ())
-> CodeGenFunction r (T final) -> CodeGenFunction r ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Exp acc -> Exp final) -> T acc -> CodeGenFunction r (T final)
forall ae am be bm r.
(Aggregate ae am, Aggregate be bm) =>
(ae -> be) -> am -> CodeGenFunction r bm
Expr.unliftM1 Exp acc -> Exp final
final T acc
accExit

mapAccumL ::
   (Shape.C sh, Marshal.C sh,
    Shape.C n, Marshal.C n,
    MultiValue.C acc, Storable.C final,
    Storable.C x,
    Storable.C y) =>
   (Exp acc -> Exp x -> Exp (acc,y)) ->
   (Exp acc -> Exp final) ->
   Sym.Array sh acc -> Sym.Array (sh, n) x ->
   (LLVM.Value (MarshalPtr sh), LLVM.Value (Ptr final)) ->
   (LLVM.Value (MarshalPtr (sh,n)), LLVM.Value (Ptr y)) ->
   LLVM.CodeGenFunction r ()
mapAccumL :: forall sh n acc final x y r.
(C sh, C sh, C n, C n, C acc, C final, C x, C y) =>
(Exp acc -> Exp x -> Exp (acc, y))
-> (Exp acc -> Exp final)
-> Array sh acc
-> Array (sh, n) x
-> (Value (MarshalPtr sh), Value (Ptr final))
-> (Value (MarshalPtr (sh, n)), Value (Ptr y))
-> CodeGenFunction r ()
mapAccumL Exp acc -> Exp x -> Exp (acc, y)
f Exp acc -> Exp final
final (Sym.Array Exp sh
_ forall r. Val (Index sh) -> Code r acc
initCode) (Sym.Array Exp (sh, n)
esh forall r. Val (Index (sh, n)) -> Code r x
code)
      (Value (MarshalPtr sh)
_, Value (Ptr final)
accPtr) (Value (MarshalPtr (sh, n))
sptr, Value (Ptr y)
yPtr) = 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 r (T (sh, n)) -> CodeGenFunction r (T sh, T n)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp (sh, n)
-> Value (MarshalPtr (sh, n)) -> CodeGenFunction r (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 (MarshalPtr (sh, n))
sptr
   let step :: Val (Index sh)
-> (Value (Ptr final), Value (Ptr y))
-> CodeGenFunction r (Value (Ptr final), Value (Ptr y))
step Val (Index sh)
ix (Value (Ptr final)
accPtr0, Value (Ptr y)
yPtrStart) = do
         T acc
accInit <- Val (Index sh) -> Code r acc
forall r. Val (Index sh) -> Code r acc
initCode Val (Index sh)
ix
         (Value (Ptr y)
ptrStop, T acc
accExit) <-
            (T (Index n) -> CodeGenFunction r (T x))
-> (Exp acc -> Exp x -> Exp (acc, y))
-> T n
-> Value (Ptr y)
-> T acc
-> CodeGenFunction r (Value (Ptr y), T acc)
forall acc b sh ix r a.
(C acc, C b, C sh, Index sh ~ ix) =>
(T ix -> CodeGenFunction r (T a))
-> (Exp acc -> Exp a -> Exp (acc, b))
-> T sh
-> Value (Ptr b)
-> T acc
-> CodeGenFunction r (Value (Ptr b), T acc)
mapAccumLLoop (T (Index sh, Index n) -> CodeGenFunction r (T x)
Val (Index (sh, n)) -> CodeGenFunction r (T x)
forall r. Val (Index (sh, n)) -> Code r x
code (T (Index sh, Index n) -> CodeGenFunction r (T x))
-> (T (Index n) -> T (Index sh, Index n))
-> T (Index n)
-> CodeGenFunction r (T x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Val (Index sh) -> T (Index n) -> T (Index sh, Index n)
forall a b. T a -> T b -> T (a, b)
MultiValue.zip Val (Index sh)
ix) Exp acc -> Exp x -> Exp (acc, y)
f T n
n Value (Ptr y)
yPtrStart T acc
accInit
         Value (Ptr final)
accPtr1 <-
            (T final
 -> Value (Ptr final) -> CodeGenFunction r (Value (Ptr final)))
-> Value (Ptr final)
-> T final
-> CodeGenFunction r (Value (Ptr final))
forall a b c. (a -> b -> c) -> b -> a -> c
flip T final
-> Value (Ptr final) -> CodeGenFunction r (Value (Ptr final))
forall a ptr r.
(C a, Value (Ptr a) ~ ptr) =>
T a -> ptr -> CodeGenFunction r ptr
Storable.storeNext Value (Ptr final)
accPtr0
               (T final -> CodeGenFunction r (Value (Ptr final)))
-> CodeGenFunction r (T final)
-> CodeGenFunction r (Value (Ptr final))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Exp acc -> Exp final) -> T acc -> CodeGenFunction r (T final)
forall ae am be bm r.
(Aggregate ae am, Aggregate be bm) =>
(ae -> be) -> am -> CodeGenFunction r bm
Expr.unliftM1 Exp acc -> Exp final
final T acc
accExit
         (Value (Ptr final), Value (Ptr y))
-> CodeGenFunction r (Value (Ptr final), Value (Ptr y))
forall a. a -> CodeGenFunction r a
forall (m :: * -> *) a. Monad m => a -> m a
return (Value (Ptr final)
accPtr1, Value (Ptr y)
ptrStop)
   CodeGenFunction r (Value (Ptr final), Value (Ptr y))
-> CodeGenFunction r ()
forall (m :: * -> *) a. Monad m => m a -> m ()
void (CodeGenFunction r (Value (Ptr final), Value (Ptr y))
 -> CodeGenFunction r ())
-> CodeGenFunction r (Value (Ptr final), Value (Ptr y))
-> CodeGenFunction r ()
forall a b. (a -> b) -> a -> b
$ (Val (Index sh)
 -> (Value (Ptr final), Value (Ptr y))
 -> CodeGenFunction r (Value (Ptr final), Value (Ptr y)))
-> T sh
-> (Value (Ptr final), Value (Ptr y))
-> CodeGenFunction r (Value (Ptr final), Value (Ptr y))
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 final), Value (Ptr y))
-> CodeGenFunction r (Value (Ptr final), Value (Ptr y))
step T sh
sh (Value (Ptr final)
accPtr,Value (Ptr y)
yPtr)

foldOuterL ::
   (Shape.C sh, Marshal.C sh,
    Shape.C n, Marshal.C n,
    Storable.C a) =>
   (Exp a -> Exp b -> Exp a) ->
   Sym.Array sh a -> Sym.Array (n,sh) b ->
   LLVM.Value (MarshalPtr sh) ->
   LLVM.Value (Ptr a) ->
   LLVM.CodeGenFunction r ()
foldOuterL :: forall sh n a b r.
(C sh, C sh, C n, C n, C a) =>
(Exp a -> Exp b -> Exp a)
-> Array sh a
-> Array (n, sh) b
-> Value (MarshalPtr sh)
-> Value (Ptr a)
-> CodeGenFunction r ()
foldOuterL Exp a -> Exp b -> Exp a
f (Sym.Array Exp sh
_ forall r. Val (Index sh) -> Code r a
initCode) (Sym.Array Exp (n, sh)
esh forall r. Val (Index (n, sh)) -> Code r b
code) Value (MarshalPtr sh)
sptr Value (Ptr a)
ptr = do
   T sh
sh <- Exp sh -> Value (MarshalPtr sh) -> CodeGenFunction r (T sh)
forall sh (f :: * -> *) r.
C sh =>
f sh -> Value (Ptr (Struct sh)) -> CodeGenFunction r (T sh)
Shape.load (Exp (n, sh) -> Exp sh
forall (val :: * -> *) a b. Value val => val (a, b) -> val b
Expr.snd Exp (n, sh)
esh) Value (MarshalPtr sh)
sptr
   T n
n <- T (n, sh) -> T n
forall a b. T (a, b) -> T a
MultiValue.fst (T (n, sh) -> T n)
-> CodeGenFunction r (T (n, sh)) -> CodeGenFunction r (T n)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp (n, sh) -> forall r. CodeGenFunction r (T (n, sh))
forall a. Exp a -> forall r. CodeGenFunction r (T a)
unExp Exp (n, sh)
esh
   CodeGenFunction r (Value (Ptr a)) -> CodeGenFunction r ()
forall (m :: * -> *) a. Monad m => m a -> m ()
void (CodeGenFunction r (Value (Ptr a)) -> CodeGenFunction r ())
-> CodeGenFunction r (Value (Ptr a)) -> CodeGenFunction r ()
forall a b. (a -> b) -> a -> b
$ T sh
-> (Val (Index sh) -> CodeGenFunction r (T a))
-> Value (Ptr a)
-> CodeGenFunction r (Value (Ptr a))
forall sh ix a r.
(C sh, Index sh ~ ix, C a) =>
T sh
-> (T ix -> CodeGenFunction r (T a))
-> Value (Ptr a)
-> CodeGenFunction r (Value (Ptr a))
writeArray T sh
sh Val (Index sh) -> CodeGenFunction r (T a)
forall r. Val (Index sh) -> Code r a
initCode Value (Ptr a)
ptr

   let step :: T (Index n)
-> Val (Index sh)
-> Value (Ptr a)
-> CodeGenFunction r (Value (Ptr a))
step T (Index n)
k Val (Index sh)
ix Value (Ptr a)
ptr0 = do
         Val b
b <- Val (Index (n, sh)) -> Code r b
forall r. Val (Index (n, sh)) -> Code r b
code (Val (Index (n, sh)) -> Code r b)
-> Val (Index (n, sh)) -> Code r b
forall a b. (a -> b) -> a -> b
$ T (Index n) -> Val (Index sh) -> T (Index n, Index sh)
forall a b. T a -> T b -> T (a, b)
MultiValue.zip T (Index n)
k Val (Index sh)
ix
         T a
a0 <- Value (Ptr a) -> CodeGenFunction r (T 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)
ptr0
         T a
a1 <- (Exp a -> Exp b -> Exp a)
-> T a -> Val b -> CodeGenFunction r (T 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 b -> Exp a
f T a
a0 Val b
b
         T a -> Value (Ptr a) -> CodeGenFunction r (Value (Ptr a))
forall a ptr r.
(C a, Value (Ptr a) ~ ptr) =>
T a -> ptr -> CodeGenFunction r ptr
Storable.storeNext T a
a1 Value (Ptr a)
ptr0
   CodeGenFunction r () -> CodeGenFunction r ()
forall (m :: * -> *) a. Monad m => m a -> m ()
void (CodeGenFunction r () -> CodeGenFunction r ())
-> CodeGenFunction r () -> CodeGenFunction r ()
forall a b. (a -> b) -> a -> b
$ (T (Index n) -> () -> CodeGenFunction r ())
-> T n -> () -> CodeGenFunction r ()
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)
k () -> CodeGenFunction r (Value (Ptr a)) -> CodeGenFunction r ()
forall (m :: * -> *) a. Monad m => m a -> m ()
void (CodeGenFunction r (Value (Ptr a)) -> CodeGenFunction r ())
-> CodeGenFunction r (Value (Ptr a)) -> CodeGenFunction r ()
forall a b. (a -> b) -> a -> b
$ (Val (Index sh)
 -> Value (Ptr a) -> CodeGenFunction r (Value (Ptr a)))
-> T sh -> Value (Ptr a) -> CodeGenFunction r (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 n)
-> Val (Index sh)
-> Value (Ptr a)
-> CodeGenFunction r (Value (Ptr a))
step T (Index n)
k) T sh
sh Value (Ptr a)
ptr) T n
n ()

{- |
We need a scalar Shape type @n@.
Scalar Shape types could be distinguished from other Shape types
by the fact that you can convert any Index into a Shape.
-}
mapFilter ::
   (Shape.Sequence n, Marshal.C n,
    Storable.C b) =>
   (Exp a -> Exp b) ->
   (Exp a -> Exp Bool) ->
   Sym.Array n a ->
   LLVM.Value (MarshalPtr n) ->
   LLVM.Value (Ptr b) ->
   LLVM.CodeGenFunction r (MultiValue.T n)
mapFilter :: forall n b a r.
(Sequence n, C n, C b) =>
(Exp a -> Exp b)
-> (Exp a -> Exp Bool)
-> Array n a
-> Value (MarshalPtr n)
-> Value (Ptr b)
-> CodeGenFunction r (T n)
mapFilter Exp a -> Exp b
f Exp a -> Exp Bool
p (Sym.Array Exp n
esh forall r. T (Index n) -> Code r a
code) Value (MarshalPtr n)
sptr Value (Ptr b)
ptr = do
   T n
n <- Exp n -> Value (MarshalPtr n) -> CodeGenFunction r (T n)
forall sh (f :: * -> *) r.
C sh =>
f sh -> Value (Ptr (Struct sh)) -> CodeGenFunction r (T sh)
Shape.load Exp n
esh Value (MarshalPtr n)
sptr
   let step :: T (Index n)
-> (Value (Ptr b), T (Index n))
-> CodeGenFunction r (Value (Ptr b), T (Index n))
step T (Index n)
ix (Value (Ptr b)
dstPtr,T (Index n)
dstIx) = do
         Val a
a <- T (Index n) -> Code r a
forall r. T (Index n) -> Code r a
code T (Index n)
ix
         MultiValue.Cons Repr Bool
c <- (Exp a -> Exp Bool) -> Val a -> CodeGenFunction r (T Bool)
forall ae am be bm r.
(Aggregate ae am, Aggregate be bm) =>
(ae -> be) -> am -> CodeGenFunction r bm
Expr.unliftM1 Exp a -> Exp Bool
p Val a
a
         Value Bool
-> (Value (Ptr b), T (Index n))
-> CodeGenFunction r (Value (Ptr b), T (Index n))
-> CodeGenFunction r (Value (Ptr b), T (Index n))
forall a r.
Phi a =>
Value Bool -> a -> CodeGenFunction r a -> CodeGenFunction r a
C.ifThen Repr Bool
Value Bool
c (Value (Ptr b)
dstPtr,T (Index n)
dstIx)
            ((Value (Ptr b) -> T (Index n) -> (Value (Ptr b), T (Index n)))
-> CodeGenFunction r (Value (Ptr b))
-> CodeGenFunction r (T (Index n))
-> CodeGenFunction r (Value (Ptr b), T (Index n))
forall (m :: * -> *) a b r.
Applicative m =>
(a -> b -> r) -> m a -> m b -> m r
App.lift2 (,)
               ((T b -> Value (Ptr b) -> CodeGenFunction r (Value (Ptr b)))
-> Value (Ptr b) -> T b -> CodeGenFunction r (Value (Ptr b))
forall a b c. (a -> b -> c) -> b -> a -> c
flip T b -> Value (Ptr b) -> CodeGenFunction r (Value (Ptr b))
forall a ptr r.
(C a, Value (Ptr a) ~ ptr) =>
T a -> ptr -> CodeGenFunction r ptr
Storable.storeNext Value (Ptr b)
dstPtr (T b -> CodeGenFunction r (Value (Ptr b)))
-> CodeGenFunction r (T b) -> CodeGenFunction r (Value (Ptr b))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Exp a -> Exp b) -> Val a -> CodeGenFunction r (T b)
forall ae am be bm r.
(Aggregate ae am, Aggregate be bm) =>
(ae -> be) -> am -> CodeGenFunction r bm
Expr.unliftM1 Exp a -> Exp b
f Val a
a)
               (T (Index n) -> CodeGenFunction r (T (Index n))
forall i r.
(Additive i, IntegerConstant i) =>
T i -> CodeGenFunction r (T i)
MultiValue.inc T (Index n)
dstIx))
   T (Index n) -> CodeGenFunction r (T n)
forall r. T (Index n) -> CodeGenFunction r (T n)
forall sh r.
Sequence sh =>
T (Index sh) -> CodeGenFunction r (T sh)
Shape.sequenceShapeFromIndex (T (Index n) -> CodeGenFunction r (T n))
-> ((Value (Ptr b), T (Index n)) -> T (Index n))
-> (Value (Ptr b), T (Index n))
-> CodeGenFunction r (T n)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value (Ptr b), T (Index n)) -> T (Index n)
forall a b. (a, b) -> b
snd
      ((Value (Ptr b), T (Index n)) -> CodeGenFunction r (T n))
-> CodeGenFunction r (Value (Ptr b), T (Index n))
-> CodeGenFunction r (T n)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (T (Index n)
 -> (Value (Ptr b), T (Index n))
 -> CodeGenFunction r (Value (Ptr b), T (Index n)))
-> T n
-> (Value (Ptr b), T (Index n))
-> CodeGenFunction r (Value (Ptr b), T (Index n))
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 b), T (Index n))
-> CodeGenFunction r (Value (Ptr b), T (Index n))
step T n
n (Value (Ptr b)
ptr, T (Index n)
forall a. C a => T a
MultiValue.zero)

filterOuter ::
   (Shape.Sequence n, Marshal.C n,
    Shape.C sh, Marshal.C sh,
    Storable.C a) =>
   Sym.Array n Bool ->
   Sym.Array (n,sh) a ->
   LLVM.Value (MarshalPtr (n,sh)) ->
   LLVM.Value (Ptr a) ->
   LLVM.CodeGenFunction r (MultiValue.T (n,sh))
filterOuter :: forall n sh a r.
(Sequence n, C n, C sh, C sh, C a) =>
Array n Bool
-> Array (n, sh) a
-> Value (MarshalPtr (n, sh))
-> Value (Ptr a)
-> CodeGenFunction r (T (n, sh))
filterOuter (Sym.Array Exp n
_eish forall r. T (Index n) -> Code r Bool
selectCode) (Sym.Array Exp (n, sh)
esh forall r. Val (Index (n, sh)) -> Code r a
code) Value (MarshalPtr (n, sh))
sptr Value (Ptr a)
ptr = do
   (T n
n,T sh
sh) <- T (n, sh) -> (T n, T sh)
forall a b. T (a, b) -> (T a, T b)
MultiValue.unzip (T (n, sh) -> (T n, T sh))
-> CodeGenFunction r (T (n, sh)) -> CodeGenFunction r (T n, T sh)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp (n, sh)
-> Value (MarshalPtr (n, sh)) -> CodeGenFunction r (T (n, sh))
forall sh (f :: * -> *) r.
C sh =>
f sh -> Value (Ptr (Struct sh)) -> CodeGenFunction r (T sh)
Shape.load Exp (n, sh)
esh Value (MarshalPtr (n, sh))
sptr
   let step :: T (Index n)
-> (Value (Ptr a), T (Index n))
-> CodeGenFunction r (Value (Ptr a), T (Index n))
step T (Index n)
k (Value (Ptr a)
dstPtr0,T (Index n)
dstK) = do
         MultiValue.Cons Repr Bool
c <- T (Index n) -> Code r Bool
forall r. T (Index n) -> Code r Bool
selectCode T (Index n)
k
         Value Bool
-> (Value (Ptr a), T (Index n))
-> CodeGenFunction r (Value (Ptr a), T (Index n))
-> CodeGenFunction r (Value (Ptr a), T (Index n))
forall a r.
Phi a =>
Value Bool -> a -> CodeGenFunction r a -> CodeGenFunction r a
C.ifThen Repr Bool
Value Bool
c (Value (Ptr a)
dstPtr0,T (Index n)
dstK)
            (do
               Value (Ptr a)
dstPtr1 <- T sh
-> (T (Index sh) -> CodeGenFunction r (T a))
-> Value (Ptr a)
-> CodeGenFunction r (Value (Ptr a))
forall sh ix a r.
(C sh, Index sh ~ ix, C a) =>
T sh
-> (T ix -> CodeGenFunction r (T a))
-> Value (Ptr a)
-> CodeGenFunction r (Value (Ptr a))
writeArray T sh
sh (T (Index n, Index sh) -> CodeGenFunction r (T a)
Val (Index (n, sh)) -> CodeGenFunction r (T a)
forall r. Val (Index (n, sh)) -> Code r a
code (T (Index n, Index sh) -> CodeGenFunction r (T a))
-> (T (Index sh) -> T (Index n, Index sh))
-> T (Index sh)
-> CodeGenFunction r (T a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T (Index n) -> T (Index sh) -> T (Index n, Index sh)
forall a b. T a -> T b -> T (a, b)
MultiValue.zip T (Index n)
k) Value (Ptr a)
dstPtr0
               (,) Value (Ptr a)
dstPtr1 (T (Index n) -> (Value (Ptr a), T (Index n)))
-> CodeGenFunction r (T (Index n))
-> CodeGenFunction r (Value (Ptr a), T (Index n))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> T (Index n) -> CodeGenFunction r (T (Index n))
forall i r.
(Additive i, IntegerConstant i) =>
T i -> CodeGenFunction r (T i)
MultiValue.inc T (Index n)
dstK)
   T n
finalN <-
      T (Index n) -> CodeGenFunction r (T n)
forall r. T (Index n) -> CodeGenFunction r (T n)
forall sh r.
Sequence sh =>
T (Index sh) -> CodeGenFunction r (T sh)
Shape.sequenceShapeFromIndex (T (Index n) -> CodeGenFunction r (T n))
-> ((Value (Ptr a), T (Index n)) -> T (Index n))
-> (Value (Ptr a), T (Index n))
-> CodeGenFunction r (T n)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value (Ptr a), T (Index n)) -> T (Index n)
forall a b. (a, b) -> b
snd
         ((Value (Ptr a), T (Index n)) -> CodeGenFunction r (T n))
-> CodeGenFunction r (Value (Ptr a), T (Index n))
-> CodeGenFunction r (T n)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (T (Index n)
 -> (Value (Ptr a), T (Index n))
 -> CodeGenFunction r (Value (Ptr a), T (Index n)))
-> T n
-> (Value (Ptr a), T (Index n))
-> CodeGenFunction r (Value (Ptr a), T (Index n))
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 (Index n))
-> CodeGenFunction r (Value (Ptr a), T (Index n))
step T n
n (Value (Ptr a)
ptr, T (Index n)
forall a. C a => T a
MultiValue.zero)
   T (n, sh) -> CodeGenFunction r (T (n, sh))
forall a. a -> CodeGenFunction r a
forall (m :: * -> *) a. Monad m => a -> m a
return (T (n, sh) -> CodeGenFunction r (T (n, sh)))
-> T (n, sh) -> CodeGenFunction r (T (n, sh))
forall a b. (a -> b) -> a -> b
$ T n -> T sh -> T (n, sh)
forall a b. T a -> T b -> T (a, b)
MultiValue.zip T n
finalN T sh
sh


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)) ->
   LLVM.Value (MarshalPtr sh1) ->
   LLVM.Value (Ptr a) ->
   LLVM.CodeGenFunction r ()
scatterMaybe :: 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 ()
scatterMaybe Exp a -> Exp a -> Exp a
accum (Sym.Array Exp sh1
esh forall r. Val (Index sh1) -> Code r a
codeInit) (Sym.Array Exp sh0
eish forall r. Val (Index sh0) -> Code r (Maybe (ix1, a))
codeMap)
      Value (MarshalPtr sh1)
sptr Value (Ptr a)
ptr = do

   T sh1
sh <- Exp sh1 -> Value (MarshalPtr sh1) -> CodeGenFunction r (T sh1)
forall sh (f :: * -> *) r.
C sh =>
f sh -> Value (Ptr (Struct sh)) -> CodeGenFunction r (T sh)
Shape.load Exp sh1
esh Value (MarshalPtr sh1)
sptr
   CodeGenFunction r (Value (Ptr a)) -> CodeGenFunction r ()
forall (m :: * -> *) a. Monad m => m a -> m ()
void (CodeGenFunction r (Value (Ptr a)) -> CodeGenFunction r ())
-> CodeGenFunction r (Value (Ptr a)) -> CodeGenFunction r ()
forall a b. (a -> b) -> a -> b
$ T sh1
-> (T ix1 -> CodeGenFunction r (T a))
-> Value (Ptr a)
-> CodeGenFunction r (Value (Ptr a))
forall sh ix a r.
(C sh, Index sh ~ ix, C a) =>
T sh
-> (T ix -> CodeGenFunction r (T a))
-> Value (Ptr a)
-> CodeGenFunction r (Value (Ptr a))
writeArray T sh1
sh T ix1 -> CodeGenFunction r (T a)
Val (Index sh1) -> CodeGenFunction r (T a)
forall r. Val (Index sh1) -> Code r a
codeInit Value (Ptr a)
ptr

   T sh0
ish <- Exp sh0 -> forall r. CodeGenFunction r (T sh0)
forall a. Exp a -> forall r. CodeGenFunction r (T a)
unExp Exp sh0
eish
   let fill :: T ix0 -> () -> CodeGenFunction r ()
fill T ix0
ix () = do
         (MultiValue.Cons Repr Bool
c, (T ix1
jx, T a
a)) <-
            (T (ix1, a) -> (T ix1, T a))
-> (T Bool, T (ix1, a)) -> (T Bool, (T ix1, T a))
forall b c a. (b -> c) -> (a, b) -> (a, c)
mapSnd T (ix1, a) -> (T ix1, T a)
forall a b. T (a, b) -> (T a, T b)
MultiValue.unzip ((T Bool, T (ix1, a)) -> (T Bool, (T ix1, T a)))
-> (T (Maybe (ix1, a)) -> (T Bool, T (ix1, a)))
-> T (Maybe (ix1, a))
-> (T Bool, (T ix1, T a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T (Maybe (ix1, a)) -> (T Bool, T (ix1, a))
forall a. T (Maybe a) -> (T Bool, T a)
MultiValue.splitMaybe (T (Maybe (ix1, a)) -> (T Bool, (T ix1, T a)))
-> CodeGenFunction r (T (Maybe (ix1, a)))
-> CodeGenFunction r (T Bool, (T ix1, T a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Val (Index sh0) -> CodeGenFunction r (T (Maybe (ix1, a)))
forall r. Val (Index sh0) -> Code r (Maybe (ix1, a))
codeMap T ix0
Val (Index sh0)
ix
         Value Bool -> () -> CodeGenFunction r () -> CodeGenFunction r ()
forall a r.
Phi a =>
Value Bool -> a -> CodeGenFunction r a -> CodeGenFunction r a
C.ifThen Repr Bool
Value Bool
c () (CodeGenFunction r () -> CodeGenFunction r ())
-> CodeGenFunction r () -> CodeGenFunction r ()
forall a b. (a -> b) -> a -> b
$ do
            Value (Ptr a)
p <- T sh1
-> Value (Ptr a) -> T ix1 -> 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 sh1
sh Value (Ptr a)
ptr T ix1
jx
            (T a -> Value (Ptr a) -> CodeGenFunction r ())
-> Value (Ptr a) -> T a -> CodeGenFunction r ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip T a -> Value (Ptr a) -> CodeGenFunction r ()
forall r. T a -> Value (Ptr a) -> CodeGenFunction r ()
forall a r. C a => T a -> Value (Ptr a) -> CodeGenFunction r ()
Storable.store Value (Ptr a)
p
               (T a -> CodeGenFunction r ())
-> CodeGenFunction r (T a) -> CodeGenFunction r ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Exp a -> Exp a -> Exp a) -> T a -> T a -> CodeGenFunction r (T 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) -> Exp a -> Exp a -> Exp a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Exp a -> Exp a -> Exp a
accum) T a
a
               (T a -> CodeGenFunction r (T a))
-> CodeGenFunction r (T a) -> CodeGenFunction r (T a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Value (Ptr a) -> CodeGenFunction r (T 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)
p
   (T ix0 -> () -> CodeGenFunction r ())
-> T sh0 -> () -> CodeGenFunction r ()
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 sh0 ~ ix, Phi state) =>
(T ix -> state -> CodeGenFunction r state)
-> T sh0 -> state -> CodeGenFunction r state
Shape.loop T ix0 -> () -> CodeGenFunction r ()
fill T sh0
ish ()

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 (Shape.Index sh1, a) ->
   LLVM.Value (MarshalPtr sh1) ->
   LLVM.Value (Ptr a) ->
   LLVM.CodeGenFunction r ()
scatter :: 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 ()
scatter Exp a -> Exp a -> Exp a
accum (Sym.Array Exp sh1
esh forall r. Val (Index sh1) -> Code r a
codeInit) (Sym.Array Exp sh0
eish forall r. Val (Index sh0) -> Code r (Index sh1, a)
codeMap) Value (MarshalPtr sh1)
sptr Value (Ptr a)
ptr = do
   T sh1
sh <- Exp sh1 -> Value (MarshalPtr sh1) -> CodeGenFunction r (T sh1)
forall sh (f :: * -> *) r.
C sh =>
f sh -> Value (Ptr (Struct sh)) -> CodeGenFunction r (T sh)
Shape.load Exp sh1
esh Value (MarshalPtr sh1)
sptr
   CodeGenFunction r (Value (Ptr a)) -> CodeGenFunction r ()
forall (m :: * -> *) a. Monad m => m a -> m ()
void (CodeGenFunction r (Value (Ptr a)) -> CodeGenFunction r ())
-> CodeGenFunction r (Value (Ptr a)) -> CodeGenFunction r ()
forall a b. (a -> b) -> a -> b
$ T sh1
-> (T ix1 -> CodeGenFunction r (T a))
-> Value (Ptr a)
-> CodeGenFunction r (Value (Ptr a))
forall sh ix a r.
(C sh, Index sh ~ ix, C a) =>
T sh
-> (T ix -> CodeGenFunction r (T a))
-> Value (Ptr a)
-> CodeGenFunction r (Value (Ptr a))
writeArray T sh1
sh T ix1 -> CodeGenFunction r (T a)
Val (Index sh1) -> CodeGenFunction r (T a)
forall r. Val (Index sh1) -> Code r a
codeInit Value (Ptr a)
ptr

   T sh0
ish <- Exp sh0 -> forall r. CodeGenFunction r (T sh0)
forall a. Exp a -> forall r. CodeGenFunction r (T a)
unExp Exp sh0
eish
   let fill :: T ix0 -> () -> CodeGenFunction r ()
fill T ix0
ix () = do
         (T ix1
jx, T a
a) <- T (ix1, a) -> (T ix1, T a)
forall a b. T (a, b) -> (T a, T b)
MultiValue.unzip (T (ix1, a) -> (T ix1, T a))
-> CodeGenFunction r (T (ix1, a)) -> CodeGenFunction r (T ix1, T a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Val (Index sh0) -> Code r (Index sh1, a)
forall r. Val (Index sh0) -> Code r (Index sh1, a)
codeMap T ix0
Val (Index sh0)
ix
         Value (Ptr a)
p <- T sh1
-> Value (Ptr a) -> T ix1 -> 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 sh1
sh Value (Ptr a)
ptr T ix1
jx
         (T a -> Value (Ptr a) -> CodeGenFunction r ())
-> Value (Ptr a) -> T a -> CodeGenFunction r ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip T a -> Value (Ptr a) -> CodeGenFunction r ()
forall r. T a -> Value (Ptr a) -> CodeGenFunction r ()
forall a r. C a => T a -> Value (Ptr a) -> CodeGenFunction r ()
Storable.store Value (Ptr a)
p
            (T a -> CodeGenFunction r ())
-> CodeGenFunction r (T a) -> CodeGenFunction r ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Exp a -> Exp a -> Exp a) -> T a -> T a -> CodeGenFunction r (T 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) -> Exp a -> Exp a -> Exp a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Exp a -> Exp a -> Exp a
accum) T a
a
            (T a -> CodeGenFunction r (T a))
-> CodeGenFunction r (T a) -> CodeGenFunction r (T a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Value (Ptr a) -> CodeGenFunction r (T 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)
p
   (T ix0 -> () -> CodeGenFunction r ())
-> T sh0 -> () -> CodeGenFunction r ()
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 sh0 ~ ix, Phi state) =>
(T ix -> state -> CodeGenFunction r state)
-> T sh0 -> state -> CodeGenFunction r state
Shape.loop T ix0 -> () -> CodeGenFunction r ()
fill T sh0
ish ()

addDimension ::
   (Shape.C n, Marshal.C n, Shape.Index n ~ k,
    Shape.C sh, Marshal.C sh,
    Storable.C b) =>
   Exp n ->
   (Exp k -> Exp a -> Exp b) ->
   Sym.Array sh a ->
   LLVM.Value (MarshalPtr (sh,n)) ->
   LLVM.Value (Ptr b) ->
   LLVM.CodeGenFunction r ()
addDimension :: forall n k sh b a r.
(C n, C n, Index n ~ k, C sh, C sh, C b) =>
Exp n
-> (Exp k -> Exp a -> Exp b)
-> Array sh a
-> Value (MarshalPtr (sh, n))
-> Value (Ptr b)
-> CodeGenFunction r ()
addDimension Exp n
en Exp k -> Exp a -> Exp b
select (Sym.Array Exp sh
esh forall r. Val (Index sh) -> Code r a
code) Value (MarshalPtr (sh, n))
sptr Value (Ptr b)
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 r (T (sh, n)) -> CodeGenFunction r (T sh, T n)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp (sh, n)
-> Value (MarshalPtr (sh, n)) -> CodeGenFunction r (T (sh, n))
forall sh (f :: * -> *) r.
C sh =>
f sh -> Value (Ptr (Struct sh)) -> CodeGenFunction r (T sh)
Shape.load (Exp sh -> Exp n -> Exp (sh, n)
forall (val :: * -> *) a b.
Value val =>
val a -> val b -> val (a, b)
Expr.zip Exp sh
esh Exp n
en) Value (MarshalPtr (sh, n))
sptr

   let fill :: Val (Index sh)
-> Value (Ptr b) -> CodeGenFunction r (Value (Ptr b))
fill Val (Index sh)
ix Value (Ptr b)
ptr0 = do
         Val a
a <- Val (Index sh) -> Code r a
forall r. Val (Index sh) -> Code r a
code Val (Index sh)
ix
         T n
-> (T k -> CodeGenFunction r (T b))
-> Value (Ptr b)
-> CodeGenFunction r (Value (Ptr b))
forall sh ix a r.
(C sh, Index sh ~ ix, C a) =>
T sh
-> (T ix -> CodeGenFunction r (T a))
-> Value (Ptr a)
-> CodeGenFunction r (Value (Ptr a))
writeArray T n
n (\T k
k -> (Exp k -> Exp a -> Exp b)
-> T k -> Val a -> CodeGenFunction r (T b)
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 k -> Exp a -> Exp b
select T k
k Val a
a) Value (Ptr b)
ptr0
   CodeGenFunction r (Value (Ptr b)) -> CodeGenFunction r ()
forall (m :: * -> *) a. Monad m => m a -> m ()
void (CodeGenFunction r (Value (Ptr b)) -> CodeGenFunction r ())
-> CodeGenFunction r (Value (Ptr b)) -> CodeGenFunction r ()
forall a b. (a -> b) -> a -> b
$ (Val (Index sh)
 -> Value (Ptr b) -> CodeGenFunction r (Value (Ptr b)))
-> T sh -> Value (Ptr b) -> CodeGenFunction r (Value (Ptr b))
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 b) -> CodeGenFunction r (Value (Ptr b))
fill T sh
sh Value (Ptr b)
ptr