{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Data.Array.Accelerate.LLVM.Native.Execute (
executeAcc,
executeOpenAcc
) where
import Data.Array.Accelerate.Analysis.Match
import Data.Array.Accelerate.Array.Unique
import Data.Array.Accelerate.Error
import Data.Array.Accelerate.Lifetime
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Shape
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Type
import Data.Array.Accelerate.LLVM.Execute
import Data.Array.Accelerate.LLVM.Execute.Async (FutureArraysR)
import Data.Array.Accelerate.LLVM.State
import Data.Array.Accelerate.LLVM.Native.Array.Data
import Data.Array.Accelerate.LLVM.Native.Execute.Async
import Data.Array.Accelerate.LLVM.Native.Execute.Divide
import Data.Array.Accelerate.LLVM.Native.Execute.Environment ( Val )
import Data.Array.Accelerate.LLVM.Native.Execute.Marshal
import Data.Array.Accelerate.LLVM.Native.Execute.Scheduler
import Data.Array.Accelerate.LLVM.Native.Link
import Data.Array.Accelerate.LLVM.Native.Target
import qualified Data.Array.Accelerate.LLVM.Native.Debug as Debug
import Control.Concurrent ( myThreadId )
import Control.Monad.State ( gets )
import Control.Monad.Trans ( liftIO )
import Data.ByteString.Short ( ShortByteString )
import Data.IORef ( newIORef, readIORef, writeIORef )
import Data.List ( find )
import Data.Maybe ( fromMaybe )
import Data.Sequence ( Seq )
import Data.Foldable ( asum )
import System.CPUTime ( getCPUTime )
import Text.Printf ( printf )
import qualified Data.ByteString.Short.Char8 as S8
import qualified Data.Sequence as Seq
import qualified Data.DList as DL
import Prelude hiding ( map, sum, scanl, scanr, init )
import Foreign.LibFFI
import Foreign.Ptr
{-# SPECIALISE INLINE executeAcc :: ExecAcc Native a -> Par Native (FutureArraysR Native a) #-}
{-# SPECIALISE INLINE executeOpenAcc :: ExecOpenAcc Native aenv a -> Val aenv -> Par Native (FutureArraysR Native a) #-}
instance Execute Native where
{-# INLINE map #-}
{-# INLINE generate #-}
{-# INLINE transform #-}
{-# INLINE backpermute #-}
{-# INLINE fold #-}
{-# INLINE foldSeg #-}
{-# INLINE scan #-}
{-# INLINE scan' #-}
{-# INLINE permute #-}
{-# INLINE stencil1 #-}
{-# INLINE stencil2 #-}
{-# INLINE aforeign #-}
map :: Maybe (a :~: b)
-> ArrayR (Array sh a)
-> TypeR b
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Array sh a
-> Par Native (FutureR Native (Array sh b))
map = Maybe (a :~: b)
-> ArrayR (Array sh a)
-> TypeR b
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Array sh a
-> Par Native (FutureR Native (Array sh b))
forall a b sh aenv.
HasCallStack =>
Maybe (a :~: b)
-> ArrayR (Array sh a)
-> TypeR b
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Array sh a
-> Par Native (Future (Array sh b))
mapOp
generate :: ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> sh
-> Par Native (FutureR Native (Array sh e))
generate = ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> sh
-> Par Native (FutureR Native (Array sh e))
forall sh e aenv.
HasCallStack =>
ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> Par Native (Future (Array sh e))
generateOp
transform :: ArrayR (Array sh a)
-> ArrayR (Array sh' b)
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> sh'
-> Array sh a
-> Par Native (FutureR Native (Array sh' b))
transform = ArrayR (Array sh a)
-> ArrayR (Array sh' b)
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> sh'
-> Array sh a
-> Par Native (FutureR Native (Array sh' b))
forall sh a sh' b aenv.
HasCallStack =>
ArrayR (Array sh a)
-> ArrayR (Array sh' b)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh'
-> Array sh a
-> Par Native (Future (Array sh' b))
transformOp
backpermute :: ArrayR (Array sh e)
-> ShapeR sh'
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> sh'
-> Array sh e
-> Par Native (FutureR Native (Array sh' e))
backpermute = ArrayR (Array sh e)
-> ShapeR sh'
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> sh'
-> Array sh e
-> Par Native (FutureR Native (Array sh' e))
forall sh e sh' aenv.
HasCallStack =>
ArrayR (Array sh e)
-> ShapeR sh'
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh'
-> Array sh e
-> Par Native (Future (Array sh' e))
backpermuteOp
fold :: HasInitialValue
-> ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (FutureR Native (Array sh e))
fold HasInitialValue
True = ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (FutureR Native (Array sh e))
forall sh e aenv.
HasCallStack =>
ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
foldOp
fold HasInitialValue
False = ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (FutureR Native (Array sh e))
forall sh e aenv.
HasCallStack =>
ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
fold1Op
foldSeg :: IntegralType i
-> HasInitialValue
-> ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Delayed (Array (sh, Int) e)
-> Delayed (Segments i)
-> Par Native (FutureR Native (Array (sh, Int) e))
foldSeg IntegralType i
i HasInitialValue
_ = IntegralType i
-> ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Delayed (Array (sh, Int) e)
-> Delayed (Segments i)
-> Par Native (Future (Array (sh, Int) e))
forall i sh e aenv.
HasCallStack =>
IntegralType i
-> ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Delayed (Segments i)
-> Par Native (Future (Array (sh, Int) e))
foldSegOp IntegralType i
i
scan :: Direction
-> HasInitialValue
-> ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (FutureR Native (Array (sh, Int) e))
scan Direction
_ HasInitialValue
True = ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (FutureR Native (Array (sh, Int) e))
forall sh e aenv.
HasCallStack =>
ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e))
scanOp
scan Direction
_ HasInitialValue
False = ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (FutureR Native (Array (sh, Int) e))
forall sh e aenv.
HasCallStack =>
ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e))
scan1Op
scan' :: Direction
-> ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (FutureR Native (Array (sh, Int) e, Array sh e))
scan' Direction
_ = ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (FutureR Native (Array (sh, Int) e, Array sh e))
forall sh e aenv.
HasCallStack =>
ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e, Array sh e))
scan'Op
permute :: HasInitialValue
-> ArrayR (Array sh e)
-> ShapeR sh'
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Array sh' e
-> Delayed (Array sh e)
-> Par Native (FutureR Native (Array sh' e))
permute = HasInitialValue
-> ArrayR (Array sh e)
-> ShapeR sh'
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Array sh' e
-> Delayed (Array sh e)
-> Par Native (FutureR Native (Array sh' e))
forall sh e sh' aenv.
HasCallStack =>
HasInitialValue
-> ArrayR (Array sh e)
-> ShapeR sh'
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Array sh' e
-> Delayed (Array sh e)
-> Par Native (Future (Array sh' e))
permuteOp
stencil1 :: TypeR a
-> ArrayR (Array sh b)
-> sh
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Delayed (Array sh a)
-> Par Native (FutureR Native (Array sh b))
stencil1 = TypeR a
-> ArrayR (Array sh b)
-> sh
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Delayed (Array sh a)
-> Par Native (FutureR Native (Array sh b))
forall a sh b aenv.
HasCallStack =>
TypeR a
-> ArrayR (Array sh b)
-> sh
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array sh a)
-> Par Native (Future (Array sh b))
stencil1Op
stencil2 :: TypeR a
-> TypeR b
-> ArrayR (Array sh c)
-> sh
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Delayed (Array sh a)
-> Delayed (Array sh b)
-> Par Native (FutureR Native (Array sh c))
stencil2 = TypeR a
-> TypeR b
-> ArrayR (Array sh c)
-> sh
-> ExecutableR Native
-> Gamma aenv
-> ValR Native aenv
-> Delayed (Array sh a)
-> Delayed (Array sh b)
-> Par Native (FutureR Native (Array sh c))
forall aenv sh a b c.
HasCallStack =>
TypeR a
-> TypeR b
-> ArrayR (Array sh c)
-> sh
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array sh a)
-> Delayed (Array sh b)
-> Par Native (Future (Array sh c))
stencil2Op
aforeign :: String
-> ArraysR as
-> ArraysR bs
-> (as -> Par Native (FutureR Native bs))
-> as
-> Par Native (FutureR Native bs)
aforeign = String
-> ArraysR as
-> ArraysR bs
-> (as -> Par Native (FutureR Native bs))
-> as
-> Par Native (FutureR Native bs)
forall as bs.
HasCallStack =>
String
-> ArraysR as
-> ArraysR bs
-> (as -> Par Native (Future bs))
-> as
-> Par Native (Future bs)
aforeignOp
{-# INLINE simpleOp #-}
simpleOp
:: HasCallStack
=> ShortByteString
-> ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> Par Native (Future (Array sh e))
simpleOp :: ShortByteString
-> ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> Par Native (Future (Array sh e))
simpleOp ShortByteString
name ArrayR (Array sh e)
repr NativeR{..} Gamma aenv
gamma Val aenv
aenv sh
sh = do
let fun :: Function
fun = Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
name
param :: TupR (ParamR Native) (Array sh e)
param = ParamR Native (Array sh e) -> TupR (ParamR Native) (Array sh e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Array sh e) -> TupR (ParamR Native) (Array sh e))
-> ParamR Native (Array sh e) -> TupR (ParamR Native) (Array sh e)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh e) -> ParamR Native (Array sh e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array sh e)
repr
Native{LinkCache
Workers
workers :: Native -> Workers
linkCache :: Native -> LinkCache
workers :: Workers
linkCache :: LinkCache
..} <- (Native -> Native) -> Par Native Native
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Native -> Native
forall t. t -> t
llvmTarget
Future (Array sh e)
future <- Par Native (Future (Array sh e))
forall arch a.
(Async arch, HasCallStack) =>
Par arch (FutureR arch a)
new
Array sh e
result <- ArrayR (Array sh e) -> sh -> Par Native (Array sh e)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote ArrayR (Array sh e)
repr sh
sh
Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> TupR (ParamR Native) (Array sh e)
-> Array sh e
-> Maybe Action
-> Par Native ()
forall aenv sh params.
HasCallStack =>
Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOp Function
fun Gamma aenv
gamma Val aenv
aenv (ArrayR (Array sh e) -> ShapeR sh
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh e)
repr) sh
sh TupR (ParamR Native) (Array sh e)
param Array sh e
result
(Maybe Action -> Par Native ()) -> Action -> Par Native ()
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers -> Future (Array sh e) -> Array sh e -> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Array sh e)
future Array sh e
result
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
Future (Array sh e) -> Par Native (Future (Array sh e))
forall (m :: * -> *) a. Monad m => a -> m a
return Future (Array sh e)
future
{-# INLINE mapOp #-}
mapOp
:: HasCallStack
=> Maybe (a :~: b)
-> ArrayR (Array sh a)
-> TypeR b
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Array sh a
-> Par Native (Future (Array sh b))
mapOp :: Maybe (a :~: b)
-> ArrayR (Array sh a)
-> TypeR b
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Array sh a
-> Par Native (Future (Array sh b))
mapOp Maybe (a :~: b)
inplace ArrayR (Array sh a)
repr TypeR b
tp NativeR{..} Gamma aenv
gamma Val aenv
aenv Array sh a
input = do
let fun :: Function
fun = Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"map"
sh :: sh
sh = Array sh a -> sh
forall sh e. Array sh e -> sh
shape Array sh a
input
shr :: ShapeR sh
shr = ArrayR (Array sh a) -> ShapeR sh
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh a)
repr
repr' :: ArrayR (Array sh b)
repr' = ShapeR sh -> TypeR b -> ArrayR (Array sh b)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
shr TypeR b
tp
param :: TupR (ParamR Native) (Array sh b, Array sh a)
param = ParamR Native (Array sh b) -> TupR (ParamR Native) (Array sh b)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Array sh b) -> ParamR Native (Array sh b)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array sh b)
repr') TupR (ParamR Native) (Array sh b)
-> TupR (ParamR Native) (Array sh a)
-> TupR (ParamR Native) (Array sh b, Array sh a)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Array sh a) -> TupR (ParamR Native) (Array sh a)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Array sh a) -> ParamR Native (Array sh a)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array sh a)
repr)
Native{LinkCache
Workers
workers :: Workers
linkCache :: LinkCache
workers :: Native -> Workers
linkCache :: Native -> LinkCache
..} <- (Native -> Native) -> Par Native Native
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Native -> Native
forall t. t -> t
llvmTarget
Future (Array sh b)
future <- Par Native (Future (Array sh b))
forall arch a.
(Async arch, HasCallStack) =>
Par arch (FutureR arch a)
new
Array sh b
result <- case Maybe (a :~: b)
inplace of
Just a :~: b
Refl -> Array sh a -> Par Native (Array sh a)
forall (m :: * -> *) a. Monad m => a -> m a
return Array sh a
input
Maybe (a :~: b)
Nothing -> ArrayR (Array sh b) -> sh -> Par Native (Array sh b)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote ArrayR (Array sh b)
repr' sh
sh
Function
-> Gamma aenv
-> Val aenv
-> ShapeR DIM1
-> DIM1
-> TupR (ParamR Native) (Array sh b, Array sh a)
-> (Array sh b, Array sh a)
-> Maybe Action
-> Par Native ()
forall aenv sh params.
HasCallStack =>
Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOp Function
fun Gamma aenv
gamma Val aenv
aenv ShapeR DIM1
dim1 ((), ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shr sh
sh) TupR (ParamR Native) (Array sh b, Array sh a)
param (Array sh b
result, Array sh a
input)
(Maybe Action -> Par Native ()) -> Action -> Par Native ()
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers -> Future (Array sh b) -> Array sh b -> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Array sh b)
future Array sh b
result
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
Future (Array sh b) -> Par Native (Future (Array sh b))
forall (m :: * -> *) a. Monad m => a -> m a
return Future (Array sh b)
future
{-# INLINE generateOp #-}
generateOp
:: HasCallStack
=> ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> Par Native (Future (Array sh e))
generateOp :: ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> Par Native (Future (Array sh e))
generateOp = ShortByteString
-> ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> Par Native (Future (Array sh e))
forall sh e aenv.
HasCallStack =>
ShortByteString
-> ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> Par Native (Future (Array sh e))
simpleOp ShortByteString
"generate"
{-# INLINE transformOp #-}
transformOp
:: HasCallStack
=> ArrayR (Array sh a)
-> ArrayR (Array sh' b)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh'
-> Array sh a
-> Par Native (Future (Array sh' b))
transformOp :: ArrayR (Array sh a)
-> ArrayR (Array sh' b)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh'
-> Array sh a
-> Par Native (Future (Array sh' b))
transformOp ArrayR (Array sh a)
repr ArrayR (Array sh' b)
repr' NativeR{..} Gamma aenv
gamma Val aenv
aenv sh'
sh' Array sh a
input = do
let fun :: Function
fun = Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"transform"
Native{LinkCache
Workers
workers :: Workers
linkCache :: LinkCache
workers :: Native -> Workers
linkCache :: Native -> LinkCache
..} <- (Native -> Native) -> Par Native Native
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Native -> Native
forall t. t -> t
llvmTarget
Future (Array sh' b)
future <- Par Native (Future (Array sh' b))
forall arch a.
(Async arch, HasCallStack) =>
Par arch (FutureR arch a)
new
Array sh' b
result <- ArrayR (Array sh' b) -> sh' -> Par Native (Array sh' b)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote ArrayR (Array sh' b)
repr' sh'
sh'
let param :: TupR (ParamR Native) (Array sh' b, Array sh a)
param = ParamR Native (Array sh' b) -> TupR (ParamR Native) (Array sh' b)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Array sh' b) -> ParamR Native (Array sh' b)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array sh' b)
repr') TupR (ParamR Native) (Array sh' b)
-> TupR (ParamR Native) (Array sh a)
-> TupR (ParamR Native) (Array sh' b, Array sh a)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Array sh a) -> TupR (ParamR Native) (Array sh a)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Array sh a) -> ParamR Native (Array sh a)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array sh a)
repr)
Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh'
-> sh'
-> TupR (ParamR Native) (Array sh' b, Array sh a)
-> (Array sh' b, Array sh a)
-> Maybe Action
-> Par Native ()
forall aenv sh params.
HasCallStack =>
Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOp Function
fun Gamma aenv
gamma Val aenv
aenv (ArrayR (Array sh' b) -> ShapeR sh'
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh' b)
repr') sh'
sh' TupR (ParamR Native) (Array sh' b, Array sh a)
param (Array sh' b
result, Array sh a
input)
(Maybe Action -> Par Native ()) -> Action -> Par Native ()
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers -> Future (Array sh' b) -> Array sh' b -> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Array sh' b)
future Array sh' b
result
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
Future (Array sh' b) -> Par Native (Future (Array sh' b))
forall (m :: * -> *) a. Monad m => a -> m a
return Future (Array sh' b)
future
{-# INLINE backpermuteOp #-}
backpermuteOp
:: HasCallStack
=> ArrayR (Array sh e)
-> ShapeR sh'
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh'
-> Array sh e
-> Par Native (Future (Array sh' e))
backpermuteOp :: ArrayR (Array sh e)
-> ShapeR sh'
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh'
-> Array sh e
-> Par Native (Future (Array sh' e))
backpermuteOp (ArrayR ShapeR sh
shr TypeR e
tp) ShapeR sh'
shr' = ArrayR (Array sh e)
-> ArrayR (Array sh' e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh'
-> Array sh e
-> Par Native (Future (Array sh' e))
forall sh a sh' b aenv.
HasCallStack =>
ArrayR (Array sh a)
-> ArrayR (Array sh' b)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh'
-> Array sh a
-> Par Native (Future (Array sh' b))
transformOp (ShapeR sh -> TypeR e -> ArrayR (Array sh e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
shr TypeR e
tp) (ShapeR sh' -> TypeR e -> ArrayR (Array sh' e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh'
shr' TypeR e
tp)
{-# INLINE fold1Op #-}
fold1Op
:: HasCallStack
=> ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
fold1Op :: ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
fold1Op ArrayR (Array sh e)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv arr :: Delayed (Array (sh, Int) e)
arr@(Delayed (Array (sh, Int) e) -> (sh, Int)
forall sh e. Delayed (Array sh e) -> sh
delayedShape -> sh :: (sh, Int)
sh@(sh
sx, Int
sz))
= String
-> HasInitialValue
-> Par Native (Future (Array sh e))
-> Par Native (Future (Array sh e))
forall a. HasCallStack => String -> HasInitialValue -> a -> a
boundsCheck String
"empty array" (Int
sz Int -> Int -> HasInitialValue
forall a. Ord a => a -> a -> HasInitialValue
> Int
0)
(Par Native (Future (Array sh e))
-> Par Native (Future (Array sh e)))
-> Par Native (Future (Array sh e))
-> Par Native (Future (Array sh e))
forall a b. (a -> b) -> a -> b
$ case ShapeR (sh, Int) -> (sh, Int) -> Int
forall sh. ShapeR sh -> sh -> Int
size (ShapeR sh -> ShapeR (sh, Int)
forall sh1. ShapeR sh1 -> ShapeR (sh1, Int)
ShapeRsnoc (ShapeR sh -> ShapeR (sh, Int)) -> ShapeR sh -> ShapeR (sh, Int)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh e) -> ShapeR sh
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh e)
repr) (sh, Int)
sh of
Int
0 -> Array sh e -> Par Native (Future (Array sh e))
forall arch a.
(Async arch, HasCallStack) =>
a -> Par arch (FutureR arch a)
newFull (Array sh e -> Par Native (Future (Array sh e)))
-> Par Native (Array sh e) -> Par Native (Future (Array sh e))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ArrayR (Array sh e) -> sh -> Par Native (Array sh e)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote ArrayR (Array sh e)
repr sh
sx
Int
_ -> ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
forall sh e aenv.
HasCallStack =>
ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
foldCore ArrayR (Array sh e)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv Delayed (Array (sh, Int) e)
arr
{-# INLINE foldOp #-}
foldOp
:: HasCallStack
=> ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
foldOp :: ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
foldOp ArrayR (Array sh e)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv arr :: Delayed (Array (sh, Int) e)
arr@(Delayed (Array (sh, Int) e) -> (sh, Int)
forall sh e. Delayed (Array sh e) -> sh
delayedShape -> sh :: (sh, Int)
sh@(sh
sx, Int
_)) =
case ShapeR (sh, Int) -> (sh, Int) -> Int
forall sh. ShapeR sh -> sh -> Int
size (ShapeR sh -> ShapeR (sh, Int)
forall sh1. ShapeR sh1 -> ShapeR (sh1, Int)
ShapeRsnoc (ShapeR sh -> ShapeR (sh, Int)) -> ShapeR sh -> ShapeR (sh, Int)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh e) -> ShapeR sh
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh e)
repr) (sh, Int)
sh of
Int
0 -> ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> Par Native (Future (Array sh e))
forall sh e aenv.
HasCallStack =>
ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> Par Native (Future (Array sh e))
generateOp ArrayR (Array sh e)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv sh
sx
Int
_ -> ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
forall sh e aenv.
HasCallStack =>
ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
foldCore ArrayR (Array sh e)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv Delayed (Array (sh, Int) e)
arr
{-# INLINE foldCore #-}
foldCore
:: HasCallStack
=> ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
foldCore :: ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
foldCore ArrayR (Array sh e)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv Delayed (Array (sh, Int) e)
arr
| ArrayR ShapeR sh
ShapeRz TypeR e
tp <- ArrayR (Array sh e)
repr
= TypeR e
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Vector e)
-> Par Native (Future (Scalar e))
forall e aenv.
HasCallStack =>
TypeR e
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Vector e)
-> Par Native (Future (Scalar e))
foldAllOp TypeR e
tp ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv Delayed (Array (sh, Int) e)
Delayed (Vector e)
arr
| HasInitialValue
otherwise
= ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
forall sh e aenv.
HasCallStack =>
ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
foldDimOp ArrayR (Array sh e)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv Delayed (Array (sh, Int) e)
arr
{-# INLINE foldAllOp #-}
foldAllOp
:: HasCallStack
=> TypeR e
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Vector e)
-> Par Native (Future (Scalar e))
foldAllOp :: TypeR e
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Vector e)
-> Par Native (Future (Scalar e))
foldAllOp TypeR e
tp NativeR{..} Gamma aenv
gamma Val aenv
aenv Delayed (Vector e)
arr = do
Native{LinkCache
Workers
workers :: Workers
linkCache :: LinkCache
workers :: Native -> Workers
linkCache :: Native -> LinkCache
..} <- (Native -> Native) -> Par Native Native
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Native -> Native
forall t. t -> t
llvmTarget
Future (Scalar e)
future <- Par Native (Future (Scalar e))
forall arch a.
(Async arch, HasCallStack) =>
Par arch (FutureR arch a)
new
Scalar e
result <- ArrayR (Scalar e) -> () -> Par Native (Scalar e)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote (ShapeR () -> TypeR e -> ArrayR (Scalar e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR ()
dim0 TypeR e
tp) ()
let
minsize :: Int
minsize = Int
4096
splits :: Int
splits = Workers -> Int
numWorkers Workers
workers
ranges :: Seq (Int, DIM1, DIM1)
ranges = Int
-> Int
-> DIM1
-> DIM1
-> (Int -> DIM1 -> DIM1 -> (Int, DIM1, DIM1))
-> Seq (Int, DIM1, DIM1)
forall a.
Int -> Int -> DIM1 -> DIM1 -> (Int -> DIM1 -> DIM1 -> a) -> Seq a
divideWork1 Int
splits Int
minsize ((), Int
0) DIM1
sh (,,)
steps :: Int
steps = Seq (Int, DIM1, DIM1) -> Int
forall a. Seq a -> Int
Seq.length Seq (Int, DIM1, DIM1)
ranges
sh :: DIM1
sh = Delayed (Vector e) -> DIM1
forall sh e. Delayed (Array sh e) -> sh
delayedShape Delayed (Vector e)
arr
if Int
steps Int -> Int -> HasInitialValue
forall a. Ord a => a -> a -> HasInitialValue
<= Int
1
then
let param :: TupR (ParamR Native) (Scalar e, Maybe (Vector e))
param = ParamR Native (Scalar e) -> TupR (ParamR Native) (Scalar e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Scalar e) -> ParamR Native (Scalar e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray (ArrayR (Scalar e) -> ParamR Native (Scalar e))
-> ArrayR (Scalar e) -> ParamR Native (Scalar e)
forall a b. (a -> b) -> a -> b
$ ShapeR () -> TypeR e -> ArrayR (Scalar e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR ()
dim0 TypeR e
tp) TupR (ParamR Native) (Scalar e)
-> TupR (ParamR Native) (Maybe (Vector e))
-> TupR (ParamR Native) (Scalar e, Maybe (Vector e))
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Maybe (Vector e))
-> TupR (ParamR Native) (Maybe (Vector e))
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Vector e) -> ParamR Native (Maybe (Vector e))
forall arch a1. ParamR arch a1 -> ParamR arch (Maybe a1)
ParamRmaybe (ParamR Native (Vector e) -> ParamR Native (Maybe (Vector e)))
-> ParamR Native (Vector e) -> ParamR Native (Maybe (Vector e))
forall a b. (a -> b) -> a -> b
$ ArrayR (Vector e) -> ParamR Native (Vector e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray (ArrayR (Vector e) -> ParamR Native (Vector e))
-> ArrayR (Vector e) -> ParamR Native (Vector e)
forall a b. (a -> b) -> a -> b
$ ShapeR DIM1 -> TypeR e -> ArrayR (Vector e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR DIM1
dim1 TypeR e
tp)
in Seq (Int, DIM1, DIM1)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR DIM1
-> TupR (ParamR Native) (Scalar e, Maybe (Vector e))
-> (Scalar e, Maybe (Vector e))
-> Maybe Action
-> Par Native ()
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOpUsing Seq (Int, DIM1, DIM1)
ranges (Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"foldAllS") Gamma aenv
gamma Val aenv
aenv ShapeR DIM1
dim1 TupR (ParamR Native) (Scalar e, Maybe (Vector e))
param (Scalar e
result, Delayed (Vector e) -> Maybe (Vector e)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Vector e)
arr)
(Maybe Action -> Par Native ()) -> Action -> Par Native ()
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers -> Future (Scalar e) -> Scalar e -> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Scalar e)
future Scalar e
result
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
else do
let param1 :: TupR (ParamR Native) (Vector e, Maybe (Vector e))
param1 = ParamR Native (Vector e) -> TupR (ParamR Native) (Vector e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Vector e) -> ParamR Native (Vector e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray (ArrayR (Vector e) -> ParamR Native (Vector e))
-> ArrayR (Vector e) -> ParamR Native (Vector e)
forall a b. (a -> b) -> a -> b
$ ShapeR DIM1 -> TypeR e -> ArrayR (Vector e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR DIM1
dim1 TypeR e
tp) TupR (ParamR Native) (Vector e)
-> TupR (ParamR Native) (Maybe (Vector e))
-> TupR (ParamR Native) (Vector e, Maybe (Vector e))
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Maybe (Vector e))
-> TupR (ParamR Native) (Maybe (Vector e))
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Vector e) -> ParamR Native (Maybe (Vector e))
forall arch a1. ParamR arch a1 -> ParamR arch (Maybe a1)
ParamRmaybe (ParamR Native (Vector e) -> ParamR Native (Maybe (Vector e)))
-> ParamR Native (Vector e) -> ParamR Native (Maybe (Vector e))
forall a b. (a -> b) -> a -> b
$ ArrayR (Vector e) -> ParamR Native (Vector e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray (ArrayR (Vector e) -> ParamR Native (Vector e))
-> ArrayR (Vector e) -> ParamR Native (Vector e)
forall a b. (a -> b) -> a -> b
$ ShapeR DIM1 -> TypeR e -> ArrayR (Vector e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR DIM1
dim1 TypeR e
tp)
let param2 :: TupR (ParamR Native) (Vector e, Scalar e)
param2 = ParamR Native (Vector e) -> TupR (ParamR Native) (Vector e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Vector e) -> ParamR Native (Vector e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray (ArrayR (Vector e) -> ParamR Native (Vector e))
-> ArrayR (Vector e) -> ParamR Native (Vector e)
forall a b. (a -> b) -> a -> b
$ ShapeR DIM1 -> TypeR e -> ArrayR (Vector e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR DIM1
dim1 TypeR e
tp) TupR (ParamR Native) (Vector e)
-> TupR (ParamR Native) (Scalar e)
-> TupR (ParamR Native) (Vector e, Scalar e)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Scalar e) -> TupR (ParamR Native) (Scalar e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Scalar e) -> ParamR Native (Scalar e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray (ArrayR (Scalar e) -> ParamR Native (Scalar e))
-> ArrayR (Scalar e) -> ParamR Native (Scalar e)
forall a b. (a -> b) -> a -> b
$ ShapeR () -> TypeR e -> ArrayR (Scalar e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR ()
dim0 TypeR e
tp)
Vector e
tmp <- ArrayR (Vector e) -> DIM1 -> Par Native (Vector e)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote (ShapeR DIM1 -> TypeR e -> ArrayR (Vector e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR DIM1
dim1 TypeR e
tp) ((), Int
steps)
Job
job2 <- Seq (Int, DIM1, DIM1)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR DIM1
-> TupR (ParamR Native) (Vector e, Scalar e)
-> (Vector e, Scalar e)
-> Maybe Action
-> Par Native Job
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJobUsing ((Int, DIM1, DIM1) -> Seq (Int, DIM1, DIM1)
forall a. a -> Seq a
Seq.singleton (Int
0, ((), Int
0), ((), Int
steps))) (Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"foldAllP2") Gamma aenv
gamma Val aenv
aenv ShapeR DIM1
dim1 TupR (ParamR Native) (Vector e, Scalar e)
param2 (Vector e
tmp, Scalar e
result)
(Maybe Action -> Par Native Job) -> Action -> Par Native Job
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers -> Future (Scalar e) -> Scalar e -> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Scalar e)
future Scalar e
result
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
Job
job1 <- Seq (Int, DIM1, DIM1)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR DIM1
-> TupR (ParamR Native) (Vector e, Maybe (Vector e))
-> (Vector e, Maybe (Vector e))
-> Maybe Action
-> Par Native Job
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJobUsingIndex Seq (Int, DIM1, DIM1)
ranges (Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"foldAllP1") Gamma aenv
gamma Val aenv
aenv ShapeR DIM1
dim1 TupR (ParamR Native) (Vector e, Maybe (Vector e))
param1 (Vector e
tmp, Delayed (Vector e) -> Maybe (Vector e)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Vector e)
arr)
(Maybe Action -> Par Native Job) -> Action -> Par Native Job
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers -> Job -> Action
schedule Workers
workers Job
job2
Action -> Par Native ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Action -> Par Native ()) -> Action -> Par Native ()
forall a b. (a -> b) -> a -> b
$ Workers -> Job -> Action
schedule Workers
workers Job
job1
Future (Scalar e) -> Par Native (Future (Scalar e))
forall (m :: * -> *) a. Monad m => a -> m a
return Future (Scalar e)
future
{-# INLINE foldDimOp #-}
foldDimOp
:: HasCallStack
=> ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
foldDimOp :: ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array sh e))
foldDimOp ArrayR (Array sh e)
repr NativeR{..} Gamma aenv
gamma Val aenv
aenv arr :: Delayed (Array (sh, Int) e)
arr@(Delayed (Array (sh, Int) e) -> (sh, Int)
forall sh e. Delayed (Array sh e) -> sh
delayedShape -> (sh
sh, Int
_)) = do
Native{LinkCache
Workers
workers :: Workers
linkCache :: LinkCache
workers :: Native -> Workers
linkCache :: Native -> LinkCache
..} <- (Native -> Native) -> Par Native Native
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Native -> Native
forall t. t -> t
llvmTarget
Future (Array sh e)
future <- Par Native (Future (Array sh e))
forall arch a.
(Async arch, HasCallStack) =>
Par arch (FutureR arch a)
new
Array sh e
result <- ArrayR (Array sh e) -> sh -> Par Native (Array sh e)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote ArrayR (Array sh e)
repr sh
sh
let
ArrayR ShapeR sh
shr TypeR e
tp = ArrayR (Array sh e)
repr
fun :: Function
fun = Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"fold"
splits :: Int
splits = Workers -> Int
numWorkers Workers
workers
minsize :: Int
minsize = Int
1
param :: TupR (ParamR Native) (Array sh e, Maybe (Array (sh, Int) e))
param = ParamR Native (Array sh e) -> TupR (ParamR Native) (Array sh e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Array sh e) -> ParamR Native (Array sh e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array sh e)
repr) TupR (ParamR Native) (Array sh e)
-> TupR (ParamR Native) (Maybe (Array (sh, Int) e))
-> TupR (ParamR Native) (Array sh e, Maybe (Array (sh, Int) e))
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Maybe (Array (sh, Int) e))
-> TupR (ParamR Native) (Maybe (Array (sh, Int) e))
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e))
forall arch a1. ParamR arch a1 -> ParamR arch (Maybe a1)
ParamRmaybe (ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e)))
-> ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e))
forall a b. (a -> b) -> a -> b
$ ArrayR (Array (sh, Int) e) -> ParamR Native (Array (sh, Int) e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray (ArrayR (Array (sh, Int) e) -> ParamR Native (Array (sh, Int) e))
-> ArrayR (Array (sh, Int) e) -> ParamR Native (Array (sh, Int) e)
forall a b. (a -> b) -> a -> b
$ ShapeR (sh, Int) -> TypeR e -> ArrayR (Array (sh, Int) e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR (ShapeR sh -> ShapeR (sh, Int)
forall sh1. ShapeR sh1 -> ShapeR (sh1, Int)
ShapeRsnoc ShapeR sh
shr) TypeR e
tp)
Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> TupR (ParamR Native) (Array sh e, Maybe (Array (sh, Int) e))
-> (Array sh e, Maybe (Array (sh, Int) e))
-> Maybe Action
-> Par Native ()
forall aenv sh params.
Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOpWith Int
splits Int
minsize Function
fun Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr sh
sh TupR (ParamR Native) (Array sh e, Maybe (Array (sh, Int) e))
param (Array sh e
result, Delayed (Array (sh, Int) e) -> Maybe (Array (sh, Int) e)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Array (sh, Int) e)
arr)
(Maybe Action -> Par Native ()) -> Action -> Par Native ()
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers -> Future (Array sh e) -> Array sh e -> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Array sh e)
future Array sh e
result
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
Future (Array sh e) -> Par Native (Future (Array sh e))
forall (m :: * -> *) a. Monad m => a -> m a
return Future (Array sh e)
future
{-# INLINE foldSegOp #-}
foldSegOp
:: HasCallStack
=> IntegralType i
-> ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Delayed (Segments i)
-> Par Native (Future (Array (sh, Int) e))
foldSegOp :: IntegralType i
-> ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Delayed (Segments i)
-> Par Native (Future (Array (sh, Int) e))
foldSegOp IntegralType i
int ArrayR (Array (sh, Int) e)
repr NativeR{..} Gamma aenv
gamma Val aenv
aenv input :: Delayed (Array (sh, Int) e)
input@(Delayed (Array (sh, Int) e) -> (sh, Int)
forall sh e. Delayed (Array sh e) -> sh
delayedShape -> (sh
sh, Int
_)) segments :: Delayed (Segments i)
segments@(Delayed (Segments i) -> DIM1
forall sh e. Delayed (Array sh e) -> sh
delayedShape -> ((), Int
ss)) = do
Native{LinkCache
Workers
workers :: Workers
linkCache :: LinkCache
workers :: Native -> Workers
linkCache :: Native -> LinkCache
..} <- (Native -> Native) -> Par Native Native
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Native -> Native
forall t. t -> t
llvmTarget
Future (Array (sh, Int) e)
future <- Par Native (Future (Array (sh, Int) e))
forall arch a.
(Async arch, HasCallStack) =>
Par arch (FutureR arch a)
new
let
n :: Int
n = Int
ssInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1
splits :: Int
splits = Workers -> Int
numWorkers Workers
workers
minsize :: Int
minsize = Int
1
shR :: ShapeR (sh, Int)
shR = ArrayR (Array (sh, Int) e) -> ShapeR (sh, Int)
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array (sh, Int) e)
repr
segR :: ArrayR (Segments i)
segR = ShapeR DIM1 -> TypeR i -> ArrayR (Segments i)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR DIM1
dim1 (TypeR i -> ArrayR (Segments i)) -> TypeR i -> ArrayR (Segments i)
forall a b. (a -> b) -> a -> b
$ ScalarType i -> TypeR i
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ScalarType i -> TypeR i) -> ScalarType i -> TypeR i
forall a b. (a -> b) -> a -> b
$ SingleType i -> ScalarType i
forall a. SingleType a -> ScalarType a
SingleScalarType (SingleType i -> ScalarType i) -> SingleType i -> ScalarType i
forall a b. (a -> b) -> a -> b
$ NumType i -> SingleType i
forall a. NumType a -> SingleType a
NumSingleType (NumType i -> SingleType i) -> NumType i -> SingleType i
forall a b. (a -> b) -> a -> b
$ IntegralType i -> NumType i
forall a. IntegralType a -> NumType a
IntegralNumType IntegralType i
int
param :: TupR
(ParamR Native)
((Array (sh, Int) e, Maybe (Array (sh, Int) e)),
Maybe (Segments i))
param = ParamR Native (Array (sh, Int) e)
-> TupR (ParamR Native) (Array (sh, Int) e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Array (sh, Int) e) -> ParamR Native (Array (sh, Int) e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array (sh, Int) e)
repr) TupR (ParamR Native) (Array (sh, Int) e)
-> TupR (ParamR Native) (Maybe (Array (sh, Int) e))
-> TupR
(ParamR Native) (Array (sh, Int) e, Maybe (Array (sh, Int) e))
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Maybe (Array (sh, Int) e))
-> TupR (ParamR Native) (Maybe (Array (sh, Int) e))
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e))
forall arch a1. ParamR arch a1 -> ParamR arch (Maybe a1)
ParamRmaybe (ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e)))
-> ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e))
forall a b. (a -> b) -> a -> b
$ ArrayR (Array (sh, Int) e) -> ParamR Native (Array (sh, Int) e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array (sh, Int) e)
repr) TupR (ParamR Native) (Array (sh, Int) e, Maybe (Array (sh, Int) e))
-> TupR (ParamR Native) (Maybe (Segments i))
-> TupR
(ParamR Native)
((Array (sh, Int) e, Maybe (Array (sh, Int) e)),
Maybe (Segments i))
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Maybe (Segments i))
-> TupR (ParamR Native) (Maybe (Segments i))
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Segments i) -> ParamR Native (Maybe (Segments i))
forall arch a1. ParamR arch a1 -> ParamR arch (Maybe a1)
ParamRmaybe (ParamR Native (Segments i) -> ParamR Native (Maybe (Segments i)))
-> ParamR Native (Segments i) -> ParamR Native (Maybe (Segments i))
forall a b. (a -> b) -> a -> b
$ ArrayR (Segments i) -> ParamR Native (Segments i)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Segments i)
segR)
Array (sh, Int) e
result <- ArrayR (Array (sh, Int) e)
-> (sh, Int) -> Par Native (Array (sh, Int) e)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote ArrayR (Array (sh, Int) e)
repr (sh
sh, Int
n)
Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR (sh, Int)
-> (sh, Int)
-> TupR
(ParamR Native)
((Array (sh, Int) e, Maybe (Array (sh, Int) e)),
Maybe (Segments i))
-> ((Array (sh, Int) e, Maybe (Array (sh, Int) e)),
Maybe (Segments i))
-> Maybe Action
-> Par Native ()
forall aenv sh params.
Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOpWith Int
splits Int
minsize (Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"foldSegP") Gamma aenv
gamma Val aenv
aenv ShapeR (sh, Int)
shR (sh
sh, Int
n) TupR
(ParamR Native)
((Array (sh, Int) e, Maybe (Array (sh, Int) e)),
Maybe (Segments i))
param ((Array (sh, Int) e
result, Delayed (Array (sh, Int) e) -> Maybe (Array (sh, Int) e)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Array (sh, Int) e)
input), Delayed (Segments i) -> Maybe (Segments i)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Segments i)
segments)
(Maybe Action -> Par Native ()) -> Action -> Par Native ()
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers
-> Future (Array (sh, Int) e) -> Array (sh, Int) e -> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Array (sh, Int) e)
future Array (sh, Int) e
result
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
Future (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e))
forall (m :: * -> *) a. Monad m => a -> m a
return Future (Array (sh, Int) e)
future
{-# INLINE scanOp #-}
scanOp
:: HasCallStack
=> ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e))
scanOp :: ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e))
scanOp ArrayR (Array (sh, Int) e)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv arr :: Delayed (Array (sh, Int) e)
arr@(Delayed (Array (sh, Int) e) -> (sh, Int)
forall sh e. Delayed (Array sh e) -> sh
delayedShape -> (sh
sz, Int
n)) =
case Int
n of
Int
0 -> ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> (sh, Int)
-> Par Native (Future (Array (sh, Int) e))
forall sh e aenv.
HasCallStack =>
ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> Par Native (Future (Array sh e))
generateOp ArrayR (Array (sh, Int) e)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv (sh
sz, Int
1)
Int
_ -> ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Int
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e))
forall sh e aenv.
HasCallStack =>
ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Int
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e))
scanCore ArrayR (Array (sh, Int) e)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Delayed (Array (sh, Int) e)
arr
{-# INLINE scan1Op #-}
scan1Op
:: HasCallStack
=> ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e))
scan1Op :: ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e))
scan1Op ArrayR (Array (sh, Int) e)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv arr :: Delayed (Array (sh, Int) e)
arr@(Delayed (Array (sh, Int) e) -> (sh, Int)
forall sh e. Delayed (Array sh e) -> sh
delayedShape -> (sh
_, Int
n))
= String
-> HasInitialValue
-> Par Native (Future (Array (sh, Int) e))
-> Par Native (Future (Array (sh, Int) e))
forall a. HasCallStack => String -> HasInitialValue -> a -> a
boundsCheck String
"empty array" (Int
n Int -> Int -> HasInitialValue
forall a. Ord a => a -> a -> HasInitialValue
> Int
0)
(Par Native (Future (Array (sh, Int) e))
-> Par Native (Future (Array (sh, Int) e)))
-> Par Native (Future (Array (sh, Int) e))
-> Par Native (Future (Array (sh, Int) e))
forall a b. (a -> b) -> a -> b
$ ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Int
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e))
forall sh e aenv.
HasCallStack =>
ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Int
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e))
scanCore ArrayR (Array (sh, Int) e)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv Int
n Delayed (Array (sh, Int) e)
arr
{-# INLINE scanCore #-}
scanCore
:: HasCallStack
=> ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Int
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e))
scanCore :: ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Int
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e))
scanCore ArrayR (Array (sh, Int) e)
repr NativeR{..} Gamma aenv
gamma Val aenv
aenv Int
m input :: Delayed (Array (sh, Int) e)
input@(Delayed (Array (sh, Int) e) -> (sh, Int)
forall sh e. Delayed (Array sh e) -> sh
delayedShape -> (sh
sz, Int
n)) = do
Native{LinkCache
Workers
workers :: Workers
linkCache :: LinkCache
workers :: Native -> Workers
linkCache :: Native -> LinkCache
..} <- (Native -> Native) -> Par Native Native
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Native -> Native
forall t. t -> t
llvmTarget
Future (Array (sh, Int) e)
future <- Par Native (Future (Array (sh, Int) e))
forall arch a.
(Async arch, HasCallStack) =>
Par arch (FutureR arch a)
new
Array (sh, Int) e
result <- ArrayR (Array (sh, Int) e)
-> (sh, Int) -> Par Native (Array (sh, Int) e)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote ArrayR (Array (sh, Int) e)
repr (sh
sz, Int
m)
let paramA :: TupR (ParamR Native) (Array (sh, Int) e)
paramA = ParamR Native (Array (sh, Int) e)
-> TupR (ParamR Native) (Array (sh, Int) e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Array (sh, Int) e)
-> TupR (ParamR Native) (Array (sh, Int) e))
-> ParamR Native (Array (sh, Int) e)
-> TupR (ParamR Native) (Array (sh, Int) e)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array (sh, Int) e) -> ParamR Native (Array (sh, Int) e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array (sh, Int) e)
repr
param :: TupR (ParamR Native) (Array (sh, Int) e, Maybe (Array (sh, Int) e))
param = TupR (ParamR Native) (Array (sh, Int) e)
paramA TupR (ParamR Native) (Array (sh, Int) e)
-> TupR (ParamR Native) (Maybe (Array (sh, Int) e))
-> TupR
(ParamR Native) (Array (sh, Int) e, Maybe (Array (sh, Int) e))
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Maybe (Array (sh, Int) e))
-> TupR (ParamR Native) (Maybe (Array (sh, Int) e))
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e))
forall arch a1. ParamR arch a1 -> ParamR arch (Maybe a1)
ParamRmaybe (ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e)))
-> ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e))
forall a b. (a -> b) -> a -> b
$ ArrayR (Array (sh, Int) e) -> ParamR Native (Array (sh, Int) e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array (sh, Int) e)
repr)
shR :: ShapeR sh
shR = ArrayR (Array sh e) -> ShapeR sh
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape (ArrayR (Array (sh, Int) e) -> ArrayR (Array sh e)
forall sh e. ArrayR (Array (sh, Int) e) -> ArrayR (Array sh e)
reduceRank ArrayR (Array (sh, Int) e)
repr)
if ShapeR (sh, Int) -> HasInitialValue
forall sh. ShapeR sh -> HasInitialValue
isMultiDim (ShapeR (sh, Int) -> HasInitialValue)
-> ShapeR (sh, Int) -> HasInitialValue
forall a b. (a -> b) -> a -> b
$ ArrayR (Array (sh, Int) e) -> ShapeR (sh, Int)
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array (sh, Int) e)
repr
then
let
fun :: Function
fun = Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"scanS"
splits :: Int
splits = Workers -> Int
numWorkers Workers
workers
minsize :: Int
minsize = Int
1
in
Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> TupR
(ParamR Native) (Array (sh, Int) e, Maybe (Array (sh, Int) e))
-> (Array (sh, Int) e, Maybe (Array (sh, Int) e))
-> Maybe Action
-> Par Native ()
forall aenv sh params.
Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOpWith Int
splits Int
minsize Function
fun Gamma aenv
gamma Val aenv
aenv ShapeR sh
shR sh
sz TupR (ParamR Native) (Array (sh, Int) e, Maybe (Array (sh, Int) e))
param (Array (sh, Int) e
result, Delayed (Array (sh, Int) e) -> Maybe (Array (sh, Int) e)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Array (sh, Int) e)
input)
(Maybe Action -> Par Native ()) -> Action -> Par Native ()
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers
-> Future (Array (sh, Int) e) -> Array (sh, Int) e -> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Array (sh, Int) e)
future Array (sh, Int) e
result
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
else
if Int
n Int -> Int -> HasInitialValue
forall a. Ord a => a -> a -> HasInitialValue
< Int
8192
then
Seq (Int, (), ())
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR ()
-> TupR
(ParamR Native) (Array (sh, Int) e, Maybe (Array (sh, Int) e))
-> (Array (sh, Int) e, Maybe (Array (sh, Int) e))
-> Maybe Action
-> Par Native ()
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOpUsing ((Int, (), ()) -> Seq (Int, (), ())
forall a. a -> Seq a
Seq.singleton (Int
0, (), ())) (Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"scanS") Gamma aenv
gamma Val aenv
aenv ShapeR ()
dim0 TupR (ParamR Native) (Array (sh, Int) e, Maybe (Array (sh, Int) e))
param (Array (sh, Int) e
result, Delayed (Array (sh, Int) e) -> Maybe (Array (sh, Int) e)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Array (sh, Int) e)
input)
(Maybe Action -> Par Native ()) -> Action -> Par Native ()
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers
-> Future (Array (sh, Int) e) -> Array (sh, Int) e -> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Array (sh, Int) e)
future Array (sh, Int) e
result
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
else do
let
splits :: Int
splits = Workers -> Int
numWorkers Workers
workers
minsize :: Int
minsize = Int
8192
ranges :: Seq (Int, DIM1, DIM1)
ranges = ShapeR DIM1
-> Int
-> Int
-> DIM1
-> DIM1
-> (Int -> DIM1 -> DIM1 -> (Int, DIM1, DIM1))
-> Seq (Int, DIM1, DIM1)
forall sh a.
ShapeR sh
-> Int -> Int -> sh -> sh -> (Int -> sh -> sh -> a) -> Seq a
divideWork ShapeR DIM1
dim1 Int
splits Int
minsize ((), Int
0) ((), Int
n) (,,)
steps :: Int
steps = Seq (Int, DIM1, DIM1) -> Int
forall a. Seq a -> Int
Seq.length Seq (Int, DIM1, DIM1)
ranges
reprTmp :: ArrayR (Array DIM1 e)
reprTmp = ShapeR DIM1 -> TypeR e -> ArrayR (Array DIM1 e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR DIM1
dim1 (TypeR e -> ArrayR (Array DIM1 e))
-> TypeR e -> ArrayR (Array DIM1 e)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array (sh, Int) e) -> TypeR e
forall sh e. ArrayR (Array sh e) -> TypeR e
arrayRtype ArrayR (Array (sh, Int) e)
repr
paramTmp :: TupR (ParamR Native) (Array DIM1 e)
paramTmp = ParamR Native (Array DIM1 e) -> TupR (ParamR Native) (Array DIM1 e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Array DIM1 e)
-> TupR (ParamR Native) (Array DIM1 e))
-> ParamR Native (Array DIM1 e)
-> TupR (ParamR Native) (Array DIM1 e)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array DIM1 e) -> ParamR Native (Array DIM1 e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array DIM1 e)
reprTmp
param1 :: TupR
(ParamR Native)
(((Int, Array (sh, Int) e), Array DIM1 e),
Maybe (Array (sh, Int) e))
param1 = ParamR Native Int -> TupR (ParamR Native) Int
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ParamR Native Int
forall arch. ParamR arch Int
ParamRint TupR (ParamR Native) Int
-> TupR (ParamR Native) (Array (sh, Int) e)
-> TupR (ParamR Native) (Int, Array (sh, Int) e)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` TupR (ParamR Native) (Array (sh, Int) e)
paramA TupR (ParamR Native) (Int, Array (sh, Int) e)
-> TupR (ParamR Native) (Array DIM1 e)
-> TupR (ParamR Native) ((Int, Array (sh, Int) e), Array DIM1 e)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` TupR (ParamR Native) (Array DIM1 e)
paramTmp TupR (ParamR Native) ((Int, Array (sh, Int) e), Array DIM1 e)
-> TupR (ParamR Native) (Maybe (Array (sh, Int) e))
-> TupR
(ParamR Native)
(((Int, Array (sh, Int) e), Array DIM1 e),
Maybe (Array (sh, Int) e))
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Maybe (Array (sh, Int) e))
-> TupR (ParamR Native) (Maybe (Array (sh, Int) e))
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e))
forall arch a1. ParamR arch a1 -> ParamR arch (Maybe a1)
ParamRmaybe (ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e)))
-> ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e))
forall a b. (a -> b) -> a -> b
$ ArrayR (Array (sh, Int) e) -> ParamR Native (Array (sh, Int) e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array (sh, Int) e)
repr)
param3 :: TupR (ParamR Native) ((Int, Array (sh, Int) e), Array DIM1 e)
param3 = ParamR Native Int -> TupR (ParamR Native) Int
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ParamR Native Int
forall arch. ParamR arch Int
ParamRint TupR (ParamR Native) Int
-> TupR (ParamR Native) (Array (sh, Int) e)
-> TupR (ParamR Native) (Int, Array (sh, Int) e)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` TupR (ParamR Native) (Array (sh, Int) e)
paramA TupR (ParamR Native) (Int, Array (sh, Int) e)
-> TupR (ParamR Native) (Array DIM1 e)
-> TupR (ParamR Native) ((Int, Array (sh, Int) e), Array DIM1 e)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` TupR (ParamR Native) (Array DIM1 e)
paramTmp
Array DIM1 e
tmp <- ArrayR (Array DIM1 e) -> DIM1 -> Par Native (Array DIM1 e)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote (ShapeR DIM1 -> TypeR e -> ArrayR (Array DIM1 e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR DIM1
dim1 (TypeR e -> ArrayR (Array DIM1 e))
-> TypeR e -> ArrayR (Array DIM1 e)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array (sh, Int) e) -> TypeR e
forall sh e. ArrayR (Array sh e) -> TypeR e
arrayRtype ArrayR (Array (sh, Int) e)
repr) ((), Int
steps)
Job
job3 <- Seq (Int, DIM1, DIM1)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR DIM1
-> TupR (ParamR Native) ((Int, Array (sh, Int) e), Array DIM1 e)
-> ((Int, Array (sh, Int) e), Array DIM1 e)
-> Maybe Action
-> Par Native Job
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJobUsingIndex Seq (Int, DIM1, DIM1)
ranges (Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"scanP3") Gamma aenv
gamma Val aenv
aenv ShapeR DIM1
dim1 TupR (ParamR Native) ((Int, Array (sh, Int) e), Array DIM1 e)
param3 ((Int
steps, Array (sh, Int) e
result), Array DIM1 e
tmp)
(Maybe Action -> Par Native Job) -> Action -> Par Native Job
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers
-> Future (Array (sh, Int) e) -> Array (sh, Int) e -> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Array (sh, Int) e)
future Array (sh, Int) e
result
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
Job
job2 <- Seq (Int, DIM1, DIM1)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR DIM1
-> TupR (ParamR Native) (Array DIM1 e)
-> Array DIM1 e
-> Maybe Action
-> Par Native Job
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJobUsing ((Int, DIM1, DIM1) -> Seq (Int, DIM1, DIM1)
forall a. a -> Seq a
Seq.singleton (Int
0, ((), Int
0), ((), Int
steps))) (Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"scanP2") Gamma aenv
gamma Val aenv
aenv ShapeR DIM1
dim1 TupR (ParamR Native) (Array DIM1 e)
paramTmp Array DIM1 e
tmp
(Maybe Action -> Par Native Job) -> Action -> Par Native Job
forall a t. (Maybe a -> t) -> a -> t
`andThen` Workers -> Job -> Action
schedule Workers
workers Job
job3
Job
job1 <- Seq (Int, DIM1, DIM1)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR DIM1
-> TupR
(ParamR Native)
(((Int, Array (sh, Int) e), Array DIM1 e),
Maybe (Array (sh, Int) e))
-> (((Int, Array (sh, Int) e), Array DIM1 e),
Maybe (Array (sh, Int) e))
-> Maybe Action
-> Par Native Job
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJobUsingIndex Seq (Int, DIM1, DIM1)
ranges (Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"scanP1") Gamma aenv
gamma Val aenv
aenv ShapeR DIM1
dim1 TupR
(ParamR Native)
(((Int, Array (sh, Int) e), Array DIM1 e),
Maybe (Array (sh, Int) e))
param1 (((Int
steps, Array (sh, Int) e
result), Array DIM1 e
tmp), Delayed (Array (sh, Int) e) -> Maybe (Array (sh, Int) e)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Array (sh, Int) e)
input)
(Maybe Action -> Par Native Job) -> Action -> Par Native Job
forall a t. (Maybe a -> t) -> a -> t
`andThen` Workers -> Job -> Action
schedule Workers
workers Job
job2
Action -> Par Native ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Action -> Par Native ()) -> Action -> Par Native ()
forall a b. (a -> b) -> a -> b
$ Workers -> Job -> Action
schedule Workers
workers Job
job1
Future (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e))
forall (m :: * -> *) a. Monad m => a -> m a
return Future (Array (sh, Int) e)
future
{-# INLINE scan'Op #-}
scan'Op
:: HasCallStack
=> ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e, Array sh e))
scan'Op :: ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e, Array sh e))
scan'Op ArrayR (Array (sh, Int) e)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv arr :: Delayed (Array (sh, Int) e)
arr@(Delayed (Array (sh, Int) e) -> (sh, Int)
forall sh e. Delayed (Array sh e) -> sh
delayedShape -> (sh
sz, Int
n)) = do
case Int
n of
Int
0 -> do
Array (sh, Int) e
out <- ArrayR (Array (sh, Int) e)
-> (sh, Int) -> Par Native (Array (sh, Int) e)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote ArrayR (Array (sh, Int) e)
repr (sh
sz, Int
0)
Future (Array sh e)
sum <- ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> Par Native (Future (Array sh e))
forall sh e aenv.
HasCallStack =>
ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> Par Native (Future (Array sh e))
generateOp (ArrayR (Array (sh, Int) e) -> ArrayR (Array sh e)
forall sh e. ArrayR (Array (sh, Int) e) -> ArrayR (Array sh e)
reduceRank ArrayR (Array (sh, Int) e)
repr) ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv sh
sz
Future (Array (sh, Int) e, Array sh e)
future <- Par Native (Future (Array (sh, Int) e, Array sh e))
forall arch a.
(Async arch, HasCallStack) =>
Par arch (FutureR arch a)
new
Par Native () -> Par Native ()
forall arch.
(Async arch, HasCallStack) =>
Par arch () -> Par arch ()
fork (Par Native () -> Par Native ()) -> Par Native () -> Par Native ()
forall a b. (a -> b) -> a -> b
$ do Array sh e
sum' <- FutureR Native (Array sh e) -> Par Native (Array sh e)
forall arch a.
(Async arch, HasCallStack) =>
FutureR arch a -> Par arch a
get FutureR Native (Array sh e)
Future (Array sh e)
sum
FutureR Native (Array (sh, Int) e, Array sh e)
-> (Array (sh, Int) e, Array sh e) -> Par Native ()
forall arch a.
(Async arch, HasCallStack) =>
FutureR arch a -> a -> Par arch ()
put FutureR Native (Array (sh, Int) e, Array sh e)
Future (Array (sh, Int) e, Array sh e)
future (Array (sh, Int) e
out, Array sh e
sum')
Future (Array (sh, Int) e, Array sh e)
-> Par Native (Future (Array (sh, Int) e, Array sh e))
forall (m :: * -> *) a. Monad m => a -> m a
return Future (Array (sh, Int) e, Array sh e)
future
Int
_ -> ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e, Array sh e))
forall aenv sh e.
HasCallStack =>
ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e, Array sh e))
scan'Core ArrayR (Array (sh, Int) e)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv Delayed (Array (sh, Int) e)
arr
{-# INLINE scan'Core #-}
scan'Core
:: forall aenv sh e. HasCallStack
=> ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e, Array sh e))
scan'Core :: ArrayR (Array (sh, Int) e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par Native (Future (Array (sh, Int) e, Array sh e))
scan'Core ArrayR (Array (sh, Int) e)
repr NativeR{..} Gamma aenv
gamma Val aenv
aenv input :: Delayed (Array (sh, Int) e)
input@(Delayed (Array (sh, Int) e) -> (sh, Int)
forall sh e. Delayed (Array sh e) -> sh
delayedShape -> sh :: (sh, Int)
sh@(sh
sz, Int
n)) = do
let
ArrayR ShapeR sh
shR TypeR e
eR = ArrayR (Array (sh, Int) e)
repr
ShapeRsnoc ShapeR sh1
shR' = ShapeR (sh, Int)
shR
repr' :: ArrayR (Array sh e)
repr' = ShapeR sh -> TypeR e -> ArrayR (Array sh e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
shR' TypeR e
eR
paramA :: TupR (ParamR Native) (Array (sh, Int) e)
paramA = ParamR Native (Array (sh, Int) e)
-> TupR (ParamR Native) (Array (sh, Int) e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Array (sh, Int) e)
-> TupR (ParamR Native) (Array (sh, Int) e))
-> ParamR Native (Array (sh, Int) e)
-> TupR (ParamR Native) (Array (sh, Int) e)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array (sh, Int) e) -> ParamR Native (Array (sh, Int) e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array (sh, Int) e)
repr
paramA' :: TupR (ParamR Native) (Array sh e)
paramA' = ParamR Native (Array sh e) -> TupR (ParamR Native) (Array sh e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Array sh e) -> TupR (ParamR Native) (Array sh e))
-> ParamR Native (Array sh e) -> TupR (ParamR Native) (Array sh e)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh e) -> ParamR Native (Array sh e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array sh e)
repr'
Native{LinkCache
Workers
workers :: Workers
linkCache :: LinkCache
workers :: Native -> Workers
linkCache :: Native -> LinkCache
..} <- (Native -> Native) -> Par Native Native
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Native -> Native
forall t. t -> t
llvmTarget
Future (Array (sh, Int) e, Array sh e)
future <- Par Native (Future (Array (sh, Int) e, Array sh e))
forall arch a.
(Async arch, HasCallStack) =>
Par arch (FutureR arch a)
new
Array (sh, Int) e
result <- ArrayR (Array (sh, Int) e)
-> (sh, Int) -> Par Native (Array (sh, Int) e)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote ArrayR (Array (sh, Int) e)
repr (sh, Int)
sh
Array sh e
sums <- ArrayR (Array sh e) -> sh -> Par Native (Array sh e)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote ArrayR (Array sh e)
repr' sh
sz
if ShapeR (sh, Int) -> HasInitialValue
forall sh. ShapeR sh -> HasInitialValue
isMultiDim ShapeR (sh, Int)
shR
then
let fun :: Function
fun = Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"scanS"
splits :: Int
splits = Workers -> Int
numWorkers Workers
workers
minsize :: Int
minsize = Int
1
param :: TupR
(ParamR Native)
((Array (sh, Int) e, Array sh e), Maybe (Array (sh, Int) e))
param = TupR (ParamR Native) (Array (sh, Int) e)
paramA TupR (ParamR Native) (Array (sh, Int) e)
-> TupR (ParamR Native) (Array sh e)
-> TupR (ParamR Native) (Array (sh, Int) e, Array sh e)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` TupR (ParamR Native) (Array sh e)
paramA' TupR (ParamR Native) (Array (sh, Int) e, Array sh e)
-> TupR (ParamR Native) (Maybe (Array (sh, Int) e))
-> TupR
(ParamR Native)
((Array (sh, Int) e, Array sh e), Maybe (Array (sh, Int) e))
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Maybe (Array (sh, Int) e))
-> TupR (ParamR Native) (Maybe (Array (sh, Int) e))
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e))
forall arch a1. ParamR arch a1 -> ParamR arch (Maybe a1)
ParamRmaybe (ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e)))
-> ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e))
forall a b. (a -> b) -> a -> b
$ ArrayR (Array (sh, Int) e) -> ParamR Native (Array (sh, Int) e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array (sh, Int) e)
repr)
in
Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> TupR
(ParamR Native)
((Array (sh, Int) e, Array sh e), Maybe (Array (sh, Int) e))
-> ((Array (sh, Int) e, Array sh e), Maybe (Array (sh, Int) e))
-> Maybe Action
-> Par Native ()
forall aenv sh params.
Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOpWith Int
splits Int
minsize Function
fun Gamma aenv
gamma Val aenv
aenv ShapeR sh
shR' sh
sz TupR
(ParamR Native)
((Array (sh, Int) e, Array sh e), Maybe (Array (sh, Int) e))
param ((Array (sh, Int) e
result, Array sh e
sums), Delayed (Array (sh, Int) e) -> Maybe (Array (sh, Int) e)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Array (sh, Int) e)
input)
(Maybe Action -> Par Native ()) -> Action -> Par Native ()
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers
-> Future (Array (sh, Int) e, Array sh e)
-> (Array (sh, Int) e, Array sh e)
-> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Array (sh, Int) e, Array sh e)
future (Array (sh, Int) e
result, Array sh e
sums)
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
else
if Int
n Int -> Int -> HasInitialValue
forall a. Ord a => a -> a -> HasInitialValue
< Int
8192
then
let param :: TupR
(ParamR Native)
((Array (sh, Int) e, Array sh e), Maybe (Array (sh, Int) e))
param = TupR (ParamR Native) (Array (sh, Int) e)
paramA TupR (ParamR Native) (Array (sh, Int) e)
-> TupR (ParamR Native) (Array sh e)
-> TupR (ParamR Native) (Array (sh, Int) e, Array sh e)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` TupR (ParamR Native) (Array sh e)
paramA' TupR (ParamR Native) (Array (sh, Int) e, Array sh e)
-> TupR (ParamR Native) (Maybe (Array (sh, Int) e))
-> TupR
(ParamR Native)
((Array (sh, Int) e, Array sh e), Maybe (Array (sh, Int) e))
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Maybe (Array (sh, Int) e))
-> TupR (ParamR Native) (Maybe (Array (sh, Int) e))
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e))
forall arch a1. ParamR arch a1 -> ParamR arch (Maybe a1)
ParamRmaybe (ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e)))
-> ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e))
forall a b. (a -> b) -> a -> b
$ ArrayR (Array (sh, Int) e) -> ParamR Native (Array (sh, Int) e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array (sh, Int) e)
repr)
in Seq (Int, (), ())
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR ()
-> TupR
(ParamR Native)
((Array (sh, Int) e, Array sh e), Maybe (Array (sh, Int) e))
-> ((Array (sh, Int) e, Array sh e), Maybe (Array (sh, Int) e))
-> Maybe Action
-> Par Native ()
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOpUsing ((Int, (), ()) -> Seq (Int, (), ())
forall a. a -> Seq a
Seq.singleton (Int
0, (), ())) (Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"scanS") Gamma aenv
gamma Val aenv
aenv ShapeR ()
dim0 TupR
(ParamR Native)
((Array (sh, Int) e, Array sh e), Maybe (Array (sh, Int) e))
param ((Array (sh, Int) e
result, Array sh e
sums), Delayed (Array (sh, Int) e) -> Maybe (Array (sh, Int) e)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Array (sh, Int) e)
input)
(Maybe Action -> Par Native ()) -> Action -> Par Native ()
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers
-> Future (Array (sh, Int) e, Array sh e)
-> (Array (sh, Int) e, Array sh e)
-> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Array (sh, Int) e, Array sh e)
future (Array (sh, Int) e
result, Array sh e
sums)
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
else do
let
splits :: Int
splits = Workers -> Int
numWorkers Workers
workers
minsize :: Int
minsize = Int
8192
ranges :: Seq (Int, DIM1, DIM1)
ranges = Int
-> Int
-> DIM1
-> DIM1
-> (Int -> DIM1 -> DIM1 -> (Int, DIM1, DIM1))
-> Seq (Int, DIM1, DIM1)
forall a.
Int -> Int -> DIM1 -> DIM1 -> (Int -> DIM1 -> DIM1 -> a) -> Seq a
divideWork1 Int
splits Int
minsize ((), Int
0) ((), Int
n) (,,)
steps :: Int
steps = Seq (Int, DIM1, DIM1) -> Int
forall a. Seq a -> Int
Seq.length Seq (Int, DIM1, DIM1)
ranges
reprTmp :: ArrayR (Array DIM1 e)
reprTmp = ShapeR DIM1 -> TypeR e -> ArrayR (Array DIM1 e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR DIM1
dim1 TypeR e
eR
paramTmp :: TupR (ParamR Native) (Array DIM1 e)
paramTmp = ParamR Native (Array DIM1 e) -> TupR (ParamR Native) (Array DIM1 e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Array DIM1 e)
-> TupR (ParamR Native) (Array DIM1 e))
-> ParamR Native (Array DIM1 e)
-> TupR (ParamR Native) (Array DIM1 e)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array DIM1 e) -> ParamR Native (Array DIM1 e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array DIM1 e)
reprTmp
param1 :: TupR
(ParamR Native)
(((Int, Array (sh, Int) e), Array DIM1 e),
Maybe (Array (sh, Int) e))
param1 = ParamR Native Int -> TupR (ParamR Native) Int
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ParamR Native Int
forall arch. ParamR arch Int
ParamRint TupR (ParamR Native) Int
-> TupR (ParamR Native) (Array (sh, Int) e)
-> TupR (ParamR Native) (Int, Array (sh, Int) e)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` TupR (ParamR Native) (Array (sh, Int) e)
paramA TupR (ParamR Native) (Int, Array (sh, Int) e)
-> TupR (ParamR Native) (Array DIM1 e)
-> TupR (ParamR Native) ((Int, Array (sh, Int) e), Array DIM1 e)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` TupR (ParamR Native) (Array DIM1 e)
paramTmp TupR (ParamR Native) ((Int, Array (sh, Int) e), Array DIM1 e)
-> TupR (ParamR Native) (Maybe (Array (sh, Int) e))
-> TupR
(ParamR Native)
(((Int, Array (sh, Int) e), Array DIM1 e),
Maybe (Array (sh, Int) e))
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Maybe (Array (sh, Int) e))
-> TupR (ParamR Native) (Maybe (Array (sh, Int) e))
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e))
forall arch a1. ParamR arch a1 -> ParamR arch (Maybe a1)
ParamRmaybe (ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e)))
-> ParamR Native (Array (sh, Int) e)
-> ParamR Native (Maybe (Array (sh, Int) e))
forall a b. (a -> b) -> a -> b
$ ArrayR (Array (sh, Int) e) -> ParamR Native (Array (sh, Int) e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array (sh, Int) e)
repr)
param2 :: TupR (ParamR Native) (Array sh e, Array DIM1 e)
param2 = TupR (ParamR Native) (Array sh e)
paramA' TupR (ParamR Native) (Array sh e)
-> TupR (ParamR Native) (Array DIM1 e)
-> TupR (ParamR Native) (Array sh e, Array DIM1 e)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` TupR (ParamR Native) (Array DIM1 e)
paramTmp
param3 :: TupR (ParamR Native) ((Int, Array (sh, Int) e), Array DIM1 e)
param3 = ParamR Native Int -> TupR (ParamR Native) Int
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ParamR Native Int
forall arch. ParamR arch Int
ParamRint TupR (ParamR Native) Int
-> TupR (ParamR Native) (Array (sh, Int) e)
-> TupR (ParamR Native) (Int, Array (sh, Int) e)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` TupR (ParamR Native) (Array (sh, Int) e)
paramA TupR (ParamR Native) (Int, Array (sh, Int) e)
-> TupR (ParamR Native) (Array DIM1 e)
-> TupR (ParamR Native) ((Int, Array (sh, Int) e), Array DIM1 e)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` TupR (ParamR Native) (Array DIM1 e)
paramTmp
Array DIM1 e
tmp <- ArrayR (Array DIM1 e) -> DIM1 -> Par Native (Array DIM1 e)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote ArrayR (Array DIM1 e)
reprTmp ((), Int
steps)
Job
job3 <- Seq (Int, DIM1, DIM1)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR DIM1
-> TupR (ParamR Native) ((Int, Array (sh, Int) e), Array DIM1 e)
-> ((Int, Array (sh, Int) e), Array DIM1 e)
-> Maybe Action
-> Par Native Job
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJobUsingIndex Seq (Int, DIM1, DIM1)
ranges (Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"scanP3") Gamma aenv
gamma Val aenv
aenv ShapeR DIM1
dim1 TupR (ParamR Native) ((Int, Array (sh, Int) e), Array DIM1 e)
param3 ((Int
steps, Array (sh, Int) e
result), Array DIM1 e
tmp)
(Maybe Action -> Par Native Job) -> Action -> Par Native Job
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers
-> Future (Array (sh, Int) e, Array sh e)
-> (Array (sh, Int) e, Array sh e)
-> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Array (sh, Int) e, Array sh e)
future (Array (sh, Int) e
result, Array sh e
sums)
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
Job
job2 <- Seq (Int, DIM1, DIM1)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR DIM1
-> TupR (ParamR Native) (Array sh e, Array DIM1 e)
-> (Array sh e, Array DIM1 e)
-> Maybe Action
-> Par Native Job
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJobUsing ((Int, DIM1, DIM1) -> Seq (Int, DIM1, DIM1)
forall a. a -> Seq a
Seq.singleton (Int
0, ((), Int
0), ((), Int
steps))) (Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"scanP2") Gamma aenv
gamma Val aenv
aenv ShapeR DIM1
dim1 TupR (ParamR Native) (Array sh e, Array DIM1 e)
param2 (Array sh e
sums, Array DIM1 e
tmp)
(Maybe Action -> Par Native Job) -> Action -> Par Native Job
forall a t. (Maybe a -> t) -> a -> t
`andThen` Workers -> Job -> Action
schedule Workers
workers Job
job3
Job
job1 <- Seq (Int, DIM1, DIM1)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR DIM1
-> TupR
(ParamR Native)
(((Int, Array (sh, Int) e), Array DIM1 e),
Maybe (Array (sh, Int) e))
-> (((Int, Array (sh, Int) e), Array DIM1 e),
Maybe (Array (sh, Int) e))
-> Maybe Action
-> Par Native Job
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJobUsingIndex Seq (Int, DIM1, DIM1)
ranges (Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"scanP1") Gamma aenv
gamma Val aenv
aenv ShapeR DIM1
dim1 TupR
(ParamR Native)
(((Int, Array (sh, Int) e), Array DIM1 e),
Maybe (Array (sh, Int) e))
param1 (((Int
steps, Array (sh, Int) e
result), Array DIM1 e
tmp), Delayed (Array (sh, Int) e) -> Maybe (Array (sh, Int) e)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Array (sh, Int) e)
input)
(Maybe Action -> Par Native Job) -> Action -> Par Native Job
forall a t. (Maybe a -> t) -> a -> t
`andThen` Workers -> Job -> Action
schedule Workers
workers Job
job2
Action -> Par Native ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Action -> Par Native ()) -> Action -> Par Native ()
forall a b. (a -> b) -> a -> b
$ Workers -> Job -> Action
schedule Workers
workers Job
job1
Future (Array (sh, Int) e, Array sh e)
-> Par Native (Future (Array (sh, Int) e, Array sh e))
forall (m :: * -> *) a. Monad m => a -> m a
return Future (Array (sh, Int) e, Array sh e)
future
isMultiDim :: ShapeR sh -> Bool
isMultiDim :: ShapeR sh -> HasInitialValue
isMultiDim (ShapeRsnoc ShapeR sh1
ShapeRz) = HasInitialValue
False
isMultiDim ShapeR sh
_ = HasInitialValue
True
{-# INLINE permuteOp #-}
permuteOp
:: forall sh e sh' aenv. HasCallStack
=> Bool
-> ArrayR (Array sh e)
-> ShapeR sh'
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Array sh' e
-> Delayed (Array sh e)
-> Par Native (Future (Array sh' e))
permuteOp :: HasInitialValue
-> ArrayR (Array sh e)
-> ShapeR sh'
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Array sh' e
-> Delayed (Array sh e)
-> Par Native (Future (Array sh' e))
permuteOp HasInitialValue
inplace ArrayR (Array sh e)
repr ShapeR sh'
shr' NativeR{..} Gamma aenv
gamma Val aenv
aenv defaults :: Array sh' e
defaults@(Array sh' e -> sh'
forall sh e. Array sh e -> sh
shape -> sh'
shOut) input :: Delayed (Array sh e)
input@(Delayed (Array sh e) -> sh
forall sh e. Delayed (Array sh e) -> sh
delayedShape -> sh
shIn) = do
let
ArrayR ShapeR sh
shr TypeR e
tp = ArrayR (Array sh e)
repr
repr' :: ArrayR (Array sh' e)
repr' = ShapeR sh' -> TypeR e -> ArrayR (Array sh' e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh'
shr' TypeR e
tp
Native{LinkCache
Workers
workers :: Workers
linkCache :: LinkCache
workers :: Native -> Workers
linkCache :: Native -> LinkCache
..} <- (Native -> Native) -> Par Native Native
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Native -> Native
forall t. t -> t
llvmTarget
Future (Array sh' e)
future <- Par Native (Future (Array sh' e))
forall arch a.
(Async arch, HasCallStack) =>
Par arch (FutureR arch a)
new
Array sh' e
result <- if HasInitialValue
inplace
then Flag
-> String -> Par Native (Array sh' e) -> Par Native (Array sh' e)
forall a. Flag -> String -> a -> a
Debug.trace Flag
Debug.dump_exec String
"exec: permute/inplace" (Par Native (Array sh' e) -> Par Native (Array sh' e))
-> Par Native (Array sh' e) -> Par Native (Array sh' e)
forall a b. (a -> b) -> a -> b
$ Array sh' e -> Par Native (Array sh' e)
forall (m :: * -> *) a. Monad m => a -> m a
return Array sh' e
defaults
else Flag
-> (Double -> Double -> String)
-> Par Native (Array sh' e)
-> Par Native (Array sh' e)
forall (m :: * -> *) a.
MonadIO m =>
Flag -> (Double -> Double -> String) -> m a -> m a
Debug.timed Flag
Debug.dump_exec (\Double
wall Double
cpu -> String
"exec: permute/clone " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Double -> Double -> String
Debug.elapsedS Double
wall Double
cpu) (Par Native (Array sh' e) -> Par Native (Array sh' e))
-> Par Native (Array sh' e) -> Par Native (Array sh' e)
forall a b. (a -> b) -> a -> b
$ LLVM Native (Array sh' e) -> Par Native (Array sh' e)
forall arch a.
(Async arch, HasCallStack) =>
LLVM arch a -> Par arch a
liftPar (ArrayR (Array sh' e) -> Array sh' e -> LLVM Native (Array sh' e)
forall sh e.
ArrayR (Array sh e) -> Array sh e -> LLVM Native (Array sh e)
cloneArray ArrayR (Array sh' e)
repr' Array sh' e
defaults)
let
splits :: Int
splits = Workers -> Int
numWorkers Workers
workers
minsize :: Int
minsize = case ShapeR sh
shr of
ShapeRsnoc ShapeR sh1
ShapeRz -> Int
4096
ShapeRsnoc (ShapeRsnoc ShapeR sh1
ShapeRz) -> Int
64
ShapeR sh
_ -> Int
16
ranges :: Seq (Int, sh, sh)
ranges = ShapeR sh
-> Int
-> Int
-> sh
-> sh
-> (Int -> sh -> sh -> (Int, sh, sh))
-> Seq (Int, sh, sh)
forall sh a.
ShapeR sh
-> Int -> Int -> sh -> sh -> (Int -> sh -> sh -> a) -> Seq a
divideWork ShapeR sh
shr Int
splits Int
minsize (ShapeR sh -> sh
forall sh. ShapeR sh -> sh
empty ShapeR sh
shr) sh
shIn (,,)
steps :: Int
steps = Seq (Int, sh, sh) -> Int
forall a. Seq a -> Int
Seq.length Seq (Int, sh, sh)
ranges
paramR :: TupR (ParamR Native) (Array sh' e, Maybe (Array sh e))
paramR = ParamR Native (Array sh' e) -> TupR (ParamR Native) (Array sh' e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Array sh' e) -> ParamR Native (Array sh' e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array sh' e)
repr') TupR (ParamR Native) (Array sh' e)
-> TupR (ParamR Native) (Maybe (Array sh e))
-> TupR (ParamR Native) (Array sh' e, Maybe (Array sh e))
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Maybe (Array sh e))
-> TupR (ParamR Native) (Maybe (Array sh e))
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Array sh e) -> ParamR Native (Maybe (Array sh e))
forall arch a1. ParamR arch a1 -> ParamR arch (Maybe a1)
ParamRmaybe (ParamR Native (Array sh e) -> ParamR Native (Maybe (Array sh e)))
-> ParamR Native (Array sh e) -> ParamR Native (Maybe (Array sh e))
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh e) -> ParamR Native (Array sh e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array sh e)
repr)
if Int
steps Int -> Int -> HasInitialValue
forall a. Ord a => a -> a -> HasInitialValue
<= Int
1
then
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> TupR (ParamR Native) (Array sh' e, Maybe (Array sh e))
-> (Array sh' e, Maybe (Array sh e))
-> Maybe Action
-> Par Native ()
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOpUsing Seq (Int, sh, sh)
ranges (Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"permuteS") Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr TupR (ParamR Native) (Array sh' e, Maybe (Array sh e))
paramR (Array sh' e
result, Delayed (Array sh e) -> Maybe (Array sh e)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Array sh e)
input)
(Maybe Action -> Par Native ()) -> Action -> Par Native ()
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers -> Future (Array sh' e) -> Array sh' e -> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Array sh' e)
future Array sh' e
result
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
else
case ShortByteString -> Lifetime FunctionTable -> Maybe Function
lookupFunction ShortByteString
"permuteP_rmw" Lifetime FunctionTable
nativeExecutable of
Just Function
f ->
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> TupR (ParamR Native) (Array sh' e, Maybe (Array sh e))
-> (Array sh' e, Maybe (Array sh e))
-> Maybe Action
-> Par Native ()
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOpUsing Seq (Int, sh, sh)
ranges Function
f Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr TupR (ParamR Native) (Array sh' e, Maybe (Array sh e))
paramR (Array sh' e
result, Delayed (Array sh e) -> Maybe (Array sh e)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Array sh e)
input)
(Maybe Action -> Par Native ()) -> Action -> Par Native ()
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers -> Future (Array sh' e) -> Array sh' e -> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Array sh' e)
future Array sh' e
result
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
Maybe Function
Nothing -> do
let m :: Int
m = ShapeR sh' -> sh' -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh'
shr' sh'
shOut
reprBarrier :: ArrayR (Array DIM1 Word8)
reprBarrier = ShapeR DIM1 -> TypeR Word8 -> ArrayR (Array DIM1 Word8)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR DIM1
dim1 (TypeR Word8 -> ArrayR (Array DIM1 Word8))
-> TypeR Word8 -> ArrayR (Array DIM1 Word8)
forall a b. (a -> b) -> a -> b
$ ScalarType Word8 -> TypeR Word8
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ScalarType Word8
scalarTypeWord8
paramR' :: TupR
(ParamR Native)
((Array sh' e, Array DIM1 Word8), Maybe (Array sh e))
paramR' = ParamR Native (Array sh' e) -> TupR (ParamR Native) (Array sh' e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Array sh' e) -> ParamR Native (Array sh' e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array sh' e)
repr') TupR (ParamR Native) (Array sh' e)
-> TupR (ParamR Native) (Array DIM1 Word8)
-> TupR (ParamR Native) (Array sh' e, Array DIM1 Word8)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Array DIM1 Word8)
-> TupR (ParamR Native) (Array DIM1 Word8)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Array DIM1 Word8) -> ParamR Native (Array DIM1 Word8)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array DIM1 Word8)
reprBarrier) TupR (ParamR Native) (Array sh' e, Array DIM1 Word8)
-> TupR (ParamR Native) (Maybe (Array sh e))
-> TupR
(ParamR Native)
((Array sh' e, Array DIM1 Word8), Maybe (Array sh e))
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Maybe (Array sh e))
-> TupR (ParamR Native) (Maybe (Array sh e))
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Array sh e) -> ParamR Native (Maybe (Array sh e))
forall arch a1. ParamR arch a1 -> ParamR arch (Maybe a1)
ParamRmaybe (ParamR Native (Array sh e) -> ParamR Native (Maybe (Array sh e)))
-> ParamR Native (Array sh e) -> ParamR Native (Maybe (Array sh e))
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh e) -> ParamR Native (Array sh e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array sh e)
repr)
barrier :: Array DIM1 Word8
barrier@(Array DIM1
_ ArrayData Word8
adb) <- ArrayR (Array DIM1 Word8) -> DIM1 -> Par Native (Array DIM1 Word8)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote ArrayR (Array DIM1 Word8)
reprBarrier ((), Int
m) :: Par Native (Vector Word8)
Action -> Par Native ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Action -> Par Native ()) -> Action -> Par Native ()
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> Word8 -> Int -> Action
memset (UniqueArray Word8 -> Ptr Word8
forall a. UniqueArray a -> Ptr a
unsafeUniqueArrayPtr ArrayData Word8
UniqueArray Word8
adb) Word8
0 Int
m
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> TupR
(ParamR Native)
((Array sh' e, Array DIM1 Word8), Maybe (Array sh e))
-> ((Array sh' e, Array DIM1 Word8), Maybe (Array sh e))
-> Maybe Action
-> Par Native ()
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOpUsing Seq (Int, sh, sh)
ranges (Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"permuteP_mutex") Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr TupR
(ParamR Native)
((Array sh' e, Array DIM1 Word8), Maybe (Array sh e))
paramR' ((Array sh' e
result, Array DIM1 Word8
barrier), Delayed (Array sh e) -> Maybe (Array sh e)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Array sh e)
input)
(Maybe Action -> Par Native ()) -> Action -> Par Native ()
forall a t. (Maybe a -> t) -> a -> t
`andThen` do Workers -> Future (Array sh' e) -> Array sh' e -> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Array sh' e)
future Array sh' e
result
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
Future (Array sh' e) -> Par Native (Future (Array sh' e))
forall (m :: * -> *) a. Monad m => a -> m a
return Future (Array sh' e)
future
{-# INLINE stencil1Op #-}
stencil1Op
:: HasCallStack
=> TypeR a
-> ArrayR (Array sh b)
-> sh
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array sh a)
-> Par Native (Future (Array sh b))
stencil1Op :: TypeR a
-> ArrayR (Array sh b)
-> sh
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array sh a)
-> Par Native (Future (Array sh b))
stencil1Op TypeR a
tp ArrayR (Array sh b)
repr sh
halo ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv input :: Delayed (Array sh a)
input@(Delayed (Array sh a) -> sh
forall sh e. Delayed (Array sh e) -> sh
delayedShape -> sh
sh) =
ArrayR (Array sh b)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> sh
-> ParamsR Native (Maybe (Array sh a))
-> Maybe (Array sh a)
-> Par Native (Future (Array sh b))
forall aenv sh e params.
HasCallStack =>
ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> sh
-> ParamsR Native params
-> params
-> Par Native (Future (Array sh e))
stencilCore ArrayR (Array sh b)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv sh
halo sh
sh (ParamR Native (Maybe (Array sh a))
-> ParamsR Native (Maybe (Array sh a))
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Maybe (Array sh a))
-> ParamsR Native (Maybe (Array sh a)))
-> ParamR Native (Maybe (Array sh a))
-> ParamsR Native (Maybe (Array sh a))
forall a b. (a -> b) -> a -> b
$ ParamR Native (Array sh a) -> ParamR Native (Maybe (Array sh a))
forall arch a1. ParamR arch a1 -> ParamR arch (Maybe a1)
ParamRmaybe (ParamR Native (Array sh a) -> ParamR Native (Maybe (Array sh a)))
-> ParamR Native (Array sh a) -> ParamR Native (Maybe (Array sh a))
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh a) -> ParamR Native (Array sh a)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray (ArrayR (Array sh a) -> ParamR Native (Array sh a))
-> ArrayR (Array sh a) -> ParamR Native (Array sh a)
forall a b. (a -> b) -> a -> b
$ ShapeR sh -> TypeR a -> ArrayR (Array sh a)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR (ArrayR (Array sh b) -> ShapeR sh
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh b)
repr) TypeR a
tp) (Delayed (Array sh a) -> Maybe (Array sh a)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Array sh a)
input)
{-# INLINE stencil2Op #-}
stencil2Op
:: forall aenv sh a b c. HasCallStack
=> TypeR a
-> TypeR b
-> ArrayR (Array sh c)
-> sh
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array sh a)
-> Delayed (Array sh b)
-> Par Native (Future (Array sh c))
stencil2Op :: TypeR a
-> TypeR b
-> ArrayR (Array sh c)
-> sh
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> Delayed (Array sh a)
-> Delayed (Array sh b)
-> Par Native (Future (Array sh c))
stencil2Op TypeR a
t1 TypeR b
t2 ArrayR (Array sh c)
repr sh
halo ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv input1 :: Delayed (Array sh a)
input1@(Delayed (Array sh a) -> sh
forall sh e. Delayed (Array sh e) -> sh
delayedShape -> sh
sh1) input2 :: Delayed (Array sh b)
input2@(Delayed (Array sh b) -> sh
forall sh e. Delayed (Array sh e) -> sh
delayedShape -> sh
sh2) =
ArrayR (Array sh c)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> sh
-> ParamsR Native (Maybe (Array sh a), Maybe (Array sh b))
-> (Maybe (Array sh a), Maybe (Array sh b))
-> Par Native (Future (Array sh c))
forall aenv sh e params.
HasCallStack =>
ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> sh
-> ParamsR Native params
-> params
-> Par Native (Future (Array sh e))
stencilCore ArrayR (Array sh c)
repr ExecutableR Native
exe Gamma aenv
gamma Val aenv
aenv sh
halo (ShapeR sh -> sh -> sh -> sh
forall sh. ShapeR sh -> sh -> sh -> sh
intersect (ArrayR (Array sh c) -> ShapeR sh
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh c)
repr) sh
sh1 sh
sh2) (TypeR a -> ParamsR Native (Maybe (Array sh a))
forall t. TypeR t -> ParamsR Native (Maybe (Array sh t))
param TypeR a
t1 ParamsR Native (Maybe (Array sh a))
-> TupR (ParamR Native) (Maybe (Array sh b))
-> ParamsR Native (Maybe (Array sh a), Maybe (Array sh b))
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` TypeR b -> TupR (ParamR Native) (Maybe (Array sh b))
forall t. TypeR t -> ParamsR Native (Maybe (Array sh t))
param TypeR b
t2) (Delayed (Array sh a) -> Maybe (Array sh a)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Array sh a)
input1, Delayed (Array sh b) -> Maybe (Array sh b)
forall sh e. Delayed (Array sh e) -> Maybe (Array sh e)
manifest Delayed (Array sh b)
input2)
where
shr :: ShapeR sh
shr = ArrayR (Array sh c) -> ShapeR sh
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh c)
repr
param :: TypeR t -> ParamsR Native (Maybe (Array sh t))
param :: TypeR t -> ParamsR Native (Maybe (Array sh t))
param = ParamR Native (Maybe (Array sh t))
-> ParamsR Native (Maybe (Array sh t))
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ParamR Native (Maybe (Array sh t))
-> ParamsR Native (Maybe (Array sh t)))
-> (TypeR t -> ParamR Native (Maybe (Array sh t)))
-> TypeR t
-> ParamsR Native (Maybe (Array sh t))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ParamR Native (Array sh t) -> ParamR Native (Maybe (Array sh t))
forall arch a1. ParamR arch a1 -> ParamR arch (Maybe a1)
ParamRmaybe (ParamR Native (Array sh t) -> ParamR Native (Maybe (Array sh t)))
-> (TypeR t -> ParamR Native (Array sh t))
-> TypeR t
-> ParamR Native (Maybe (Array sh t))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ArrayR (Array sh t) -> ParamR Native (Array sh t)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray (ArrayR (Array sh t) -> ParamR Native (Array sh t))
-> (TypeR t -> ArrayR (Array sh t))
-> TypeR t
-> ParamR Native (Array sh t)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeR sh -> TypeR t -> ArrayR (Array sh t)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
shr
{-# INLINE stencilCore #-}
stencilCore
:: forall aenv sh e params. HasCallStack
=> ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> sh
-> ParamsR Native params
-> params
-> Par Native (Future (Array sh e))
stencilCore :: ArrayR (Array sh e)
-> ExecutableR Native
-> Gamma aenv
-> Val aenv
-> sh
-> sh
-> ParamsR Native params
-> params
-> Par Native (Future (Array sh e))
stencilCore ArrayR (Array sh e)
repr NativeR{..} Gamma aenv
gamma Val aenv
aenv sh
halo sh
sh ParamsR Native params
paramsR params
params = do
Native{LinkCache
Workers
workers :: Workers
linkCache :: LinkCache
workers :: Native -> Workers
linkCache :: Native -> LinkCache
..} <- (Native -> Native) -> Par Native Native
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Native -> Native
forall t. t -> t
llvmTarget
Future (Array sh e)
future <- Par Native (Future (Array sh e))
forall arch a.
(Async arch, HasCallStack) =>
Par arch (FutureR arch a)
new
Array sh e
result <- ArrayR (Array sh e) -> sh -> Par Native (Array sh e)
forall arch sh e.
Remote arch =>
ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
allocateRemote ArrayR (Array sh e)
repr sh
sh
let
shr :: ShapeR sh
shr = ArrayR (Array sh e) -> ShapeR sh
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh e)
repr
inside :: Function
inside = Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"stencil_inside"
border :: Function
border = Lifetime FunctionTable
nativeExecutable HasCallStack =>
Lifetime FunctionTable -> ShortByteString -> Function
Lifetime FunctionTable -> ShortByteString -> Function
!# ShortByteString
"stencil_border"
splits :: Int
splits = Workers -> Int
numWorkers Workers
workers
minsize :: Int
minsize = case ShapeR sh
shr of
ShapeRsnoc ShapeR sh1
ShapeRz -> Int
4096
ShapeRsnoc (ShapeRsnoc ShapeR sh1
ShapeRz) -> Int
64
ShapeR sh
_ -> Int
16
ins :: Seq (Int, sh, sh)
ins = ShapeR sh
-> Int
-> Int
-> sh
-> sh
-> (Int -> sh -> sh -> (Int, sh, sh))
-> Seq (Int, sh, sh)
forall sh a.
ShapeR sh
-> Int -> Int -> sh -> sh -> (Int -> sh -> sh -> a) -> Seq a
divideWork ShapeR sh
shr Int
splits Int
minsize sh
halo (sh -> sh -> sh
sub sh
sh sh
halo) (,,)
outs :: Seq (Int, sh, sh)
outs = Seq (Seq (Int, sh, sh)) -> Seq (Int, sh, sh)
forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Alternative f) =>
t (f a) -> f a
asum (Seq (Seq (Int, sh, sh)) -> Seq (Int, sh, sh))
-> (((sh, sh) -> Seq (Int, sh, sh)) -> Seq (Seq (Int, sh, sh)))
-> ((sh, sh) -> Seq (Int, sh, sh))
-> Seq (Int, sh, sh)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (((sh, sh) -> Seq (Int, sh, sh))
-> Seq (sh, sh) -> Seq (Seq (Int, sh, sh)))
-> Seq (sh, sh)
-> ((sh, sh) -> Seq (Int, sh, sh))
-> Seq (Seq (Int, sh, sh))
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((sh, sh) -> Seq (Int, sh, sh))
-> Seq (sh, sh) -> Seq (Seq (Int, sh, sh))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ShapeR sh -> sh -> sh -> Seq (sh, sh)
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Seq (sh, sh)
stencilBorders ShapeR sh
shr sh
sh sh
halo) (((sh, sh) -> Seq (Int, sh, sh)) -> Seq (Int, sh, sh))
-> ((sh, sh) -> Seq (Int, sh, sh)) -> Seq (Int, sh, sh)
forall a b. (a -> b) -> a -> b
$ \(sh
u,sh
v) -> ShapeR sh
-> Int
-> Int
-> sh
-> sh
-> (Int -> sh -> sh -> (Int, sh, sh))
-> Seq (Int, sh, sh)
forall sh a.
ShapeR sh
-> Int -> Int -> sh -> sh -> (Int -> sh -> sh -> a) -> Seq a
divideWork ShapeR sh
shr Int
splits Int
minsize sh
u sh
v (,,)
sub :: sh -> sh -> sh
sub :: sh -> sh -> sh
sub sh
a sh
b = ShapeR sh -> sh -> sh -> sh
forall sh. ShapeR sh -> sh -> sh -> sh
go ShapeR sh
shr sh
a sh
b
where
go :: ShapeR t -> t -> t -> t
go :: ShapeR t -> t -> t -> t
go ShapeR t
ShapeRz () () = ()
go (ShapeRsnoc ShapeR sh1
shr') (xa,xb) (ya,yb) = (ShapeR sh1 -> sh1 -> sh1 -> sh1
forall sh. ShapeR sh -> sh -> sh -> sh
go ShapeR sh1
shr' sh1
xa sh1
ya, Int
xb Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
yb)
paramsR' :: TupR (ParamR Native) (Array sh e, params)
paramsR' = ParamR Native (Array sh e) -> TupR (ParamR Native) (Array sh e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Array sh e) -> ParamR Native (Array sh e)
forall sh e arch. ArrayR (Array sh e) -> ParamR arch (Array sh e)
ParamRarray ArrayR (Array sh e)
repr) TupR (ParamR Native) (Array sh e)
-> ParamsR Native params
-> TupR (ParamR Native) (Array sh e, params)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamsR Native params
paramsR
Seq Action
jobsInside <- Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> TupR (ParamR Native) (Array sh e, params)
-> (Array sh e, params)
-> Par Native (Seq Action)
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Par Native (Seq Action)
mkTasksUsing Seq (Int, sh, sh)
ins Function
inside Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr TupR (ParamR Native) (Array sh e, params)
paramsR' (Array sh e
result, params
params)
Seq Action
jobsBorder <- Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> TupR (ParamR Native) (Array sh e, params)
-> (Array sh e, params)
-> Par Native (Seq Action)
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Par Native (Seq Action)
mkTasksUsing Seq (Int, sh, sh)
outs Function
border Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr TupR (ParamR Native) (Array sh e, params)
paramsR' (Array sh e
result, params
params)
let jobTasks :: Seq Action
jobTasks = Seq Action
jobsInside Seq Action -> Seq Action -> Seq Action
forall a. Seq a -> Seq a -> Seq a
Seq.>< Seq Action
jobsBorder
jobDone :: Maybe Action
jobDone = Action -> Maybe Action
forall a. a -> Maybe a
Just (Action -> Maybe Action) -> Action -> Maybe Action
forall a b. (a -> b) -> a -> b
$ do Workers -> Future (Array sh e) -> Array sh e -> Action
forall a. HasCallStack => Workers -> Future a -> a -> Action
putIO Workers
workers Future (Array sh e)
future Array sh e
result
Lifetime FunctionTable -> Action
forall a. Lifetime a -> Action
touchLifetime Lifetime FunctionTable
nativeExecutable
Action -> Par Native ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Action -> Par Native ()) -> Action -> Par Native ()
forall a b. (a -> b) -> a -> b
$ Workers -> Job -> Action
schedule Workers
workers (Job -> Action) -> IO Job -> Action
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ShortByteString -> Job -> IO Job
timed ShortByteString
"stencil" Job :: Seq Action -> Maybe Action -> Job
Job{Maybe Action
Seq Action
jobDone :: Maybe Action
jobTasks :: Seq Action
jobDone :: Maybe Action
jobTasks :: Seq Action
..}
Future (Array sh e) -> Par Native (Future (Array sh e))
forall (m :: * -> *) a. Monad m => a -> m a
return Future (Array sh e)
future
{-# INLINE stencilBorders #-}
stencilBorders
:: forall sh. HasCallStack
=> ShapeR sh
-> sh
-> sh
-> Seq (sh, sh)
stencilBorders :: ShapeR sh -> sh -> sh -> Seq (sh, sh)
stencilBorders ShapeR sh
shr sh
sh sh
halo = Int -> (Int -> (sh, sh)) -> Seq (sh, sh)
forall a. Int -> (Int -> a) -> Seq a
Seq.fromFunction (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* ShapeR sh -> Int
forall sh. ShapeR sh -> Int
rank ShapeR sh
shr) Int -> (sh, sh)
face
where
face :: Int -> (sh, sh)
face :: Int -> (sh, sh)
face Int
n = Int -> ShapeR sh -> sh -> sh -> (sh, sh)
forall t. Int -> ShapeR t -> t -> t -> (t, t)
go Int
n ShapeR sh
shr sh
sh sh
halo
go :: Int -> ShapeR t -> t -> t -> (t, t)
go :: Int -> ShapeR t -> t -> t -> (t, t)
go Int
_ ShapeR t
ShapeRz () () = ((), ())
go Int
n (ShapeRsnoc ShapeR sh1
shr') (sha, sza) (shb, szb)
= let
(sh1
sha', sh1
shb') = Int -> ShapeR sh1 -> sh1 -> sh1 -> (sh1, sh1)
forall t. Int -> ShapeR t -> t -> t -> (t, t)
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
2) ShapeR sh1
shr' sh1
sha sh1
shb
(Int
sza', Int
szb')
| Int
n Int -> Int -> HasInitialValue
forall a. Ord a => a -> a -> HasInitialValue
< Int
0 = (Int
0, Int
sza)
| Int
n Int -> Int -> HasInitialValue
forall a. Eq a => a -> a -> HasInitialValue
== Int
0 = (Int
0, Int
szb)
| Int
n Int -> Int -> HasInitialValue
forall a. Eq a => a -> a -> HasInitialValue
== Int
1 = (Int
szaInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
szb, Int
sza)
| HasInitialValue
otherwise = (Int
szb, Int
szaInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
szb)
in
((sh1
sha', Int
sza'), (sh1
shb', Int
szb'))
{-# INLINE aforeignOp #-}
aforeignOp
:: HasCallStack
=> String
-> ArraysR as
-> ArraysR bs
-> (as -> Par Native (Future bs))
-> as
-> Par Native (Future bs)
aforeignOp :: String
-> ArraysR as
-> ArraysR bs
-> (as -> Par Native (Future bs))
-> as
-> Par Native (Future bs)
aforeignOp String
name ArraysR as
_ ArraysR bs
_ as -> Par Native (Future bs)
asm as
arr = do
Double
wallBegin <- IO Double -> Par Native Double
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Double
getMonotonicTime
Future bs
result <- Flag
-> (Double -> Double -> String)
-> Par Native (Future bs)
-> Par Native (Future bs)
forall (m :: * -> *) a.
MonadIO m =>
Flag -> (Double -> Double -> String) -> m a -> m a
Debug.timed Flag
Debug.dump_exec (\Double
wall Double
cpu -> String -> String -> String -> String
forall r. PrintfType r => String -> r
printf String
"exec: %s %s" String
name (Double -> Double -> String
Debug.elapsedP Double
wall Double
cpu)) (as -> Par Native (Future bs)
asm as
arr)
Double
wallEnd <- IO Double -> Par Native Double
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Double
getMonotonicTime
Action -> Par Native ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Action -> Par Native ()) -> Action -> Par Native ()
forall a b. (a -> b) -> a -> b
$ Processor -> Double -> Action
Debug.addProcessorTime Processor
Debug.Native (Double
wallEnd Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
wallBegin)
Future bs -> Par Native (Future bs)
forall (m :: * -> *) a. Monad m => a -> m a
return Future bs
result
(!#) :: HasCallStack => Lifetime FunctionTable -> ShortByteString -> Function
!# :: Lifetime FunctionTable -> ShortByteString -> Function
(!#) Lifetime FunctionTable
exe ShortByteString
name
= Function -> Maybe Function -> Function
forall a. a -> Maybe a -> a
fromMaybe (String -> Function
forall a. HasCallStack => String -> a
internalError (String
"function not found: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ ShortByteString -> String
S8.unpack ShortByteString
name))
(Maybe Function -> Function) -> Maybe Function -> Function
forall a b. (a -> b) -> a -> b
$ ShortByteString -> Lifetime FunctionTable -> Maybe Function
lookupFunction ShortByteString
name Lifetime FunctionTable
exe
lookupFunction :: ShortByteString -> Lifetime FunctionTable -> Maybe Function
lookupFunction :: ShortByteString -> Lifetime FunctionTable -> Maybe Function
lookupFunction ShortByteString
name Lifetime FunctionTable
nativeExecutable = do
(Function -> HasInitialValue) -> [Function] -> Maybe Function
forall (t :: * -> *) a.
Foldable t =>
(a -> HasInitialValue) -> t a -> Maybe a
find (\(ShortByteString
n,FunPtr ()
_) -> ShortByteString
n ShortByteString -> ShortByteString -> HasInitialValue
forall a. Eq a => a -> a -> HasInitialValue
== ShortByteString
name) (FunctionTable -> [Function]
functionTable (Lifetime FunctionTable -> FunctionTable
forall a. Lifetime a -> a
unsafeGetValue Lifetime FunctionTable
nativeExecutable))
andThen :: (Maybe a -> t) -> a -> t
andThen :: (Maybe a -> t) -> a -> t
andThen Maybe a -> t
f a
g = Maybe a -> t
f (a -> Maybe a
forall a. a -> Maybe a
Just a
g)
delayedShape :: Delayed (Array sh e) -> sh
delayedShape :: Delayed (Array sh e) -> sh
delayedShape (Delayed sh
sh) = sh
sh
sh
delayedShape (Manifest Array sh e
a) = Array sh e -> sh
forall sh e. Array sh e -> sh
shape Array sh e
a
manifest :: Delayed (Array sh e) -> Maybe (Array sh e)
manifest :: Delayed (Array sh e) -> Maybe (Array sh e)
manifest (Manifest Array sh e
a) = Array sh e -> Maybe (Array sh e)
forall a. a -> Maybe a
Just Array sh e
a
manifest Delayed{} = Maybe (Array sh e)
forall a. Maybe a
Nothing
{-# INLINABLE scheduleOp #-}
scheduleOp
:: HasCallStack
=> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOp :: Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOp Function
fun Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr sh
sz ParamsR Native params
paramsR params
params Maybe Action
done = do
Native{LinkCache
Workers
workers :: Workers
linkCache :: LinkCache
workers :: Native -> Workers
linkCache :: Native -> LinkCache
..} <- (Native -> Native) -> Par Native Native
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Native -> Native
forall t. t -> t
llvmTarget
let
splits :: Int
splits = Workers -> Int
numWorkers Workers
workers
minsize :: Int
minsize = case ShapeR sh
shr of
ShapeRsnoc ShapeR sh1
ShapeRz -> Int
4096
ShapeRsnoc (ShapeRsnoc ShapeR sh1
ShapeRz) -> Int
64
ShapeR sh
_ -> Int
16
Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
forall aenv sh params.
Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOpWith Int
splits Int
minsize Function
fun Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr sh
sz ParamsR Native params
paramsR params
params Maybe Action
done
{-# INLINABLE scheduleOpWith #-}
scheduleOpWith
:: Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOpWith :: Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOpWith Int
splits Int
minsize Function
fun Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr sh
sz ParamsR Native params
paramsR params
params Maybe Action
done = do
Native{LinkCache
Workers
workers :: Workers
linkCache :: LinkCache
workers :: Native -> Workers
linkCache :: Native -> LinkCache
..} <- (Native -> Native) -> Par Native Native
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Native -> Native
forall t. t -> t
llvmTarget
Job
job <- Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
forall aenv sh params.
Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJob Int
splits Int
minsize Function
fun Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr (ShapeR sh -> sh
forall sh. ShapeR sh -> sh
empty ShapeR sh
shr) sh
sz ParamsR Native params
paramsR params
params Maybe Action
done
Action -> Par Native ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Action -> Par Native ()) -> Action -> Par Native ()
forall a b. (a -> b) -> a -> b
$ Workers -> Job -> Action
schedule Workers
workers Job
job
{-# INLINABLE scheduleOpUsing #-}
scheduleOpUsing
:: Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOpUsing :: Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native ()
scheduleOpUsing Seq (Int, sh, sh)
ranges Function
fun Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr ParamsR Native params
paramsR params
params Maybe Action
jobDone = do
Native{LinkCache
Workers
workers :: Workers
linkCache :: LinkCache
workers :: Native -> Workers
linkCache :: Native -> LinkCache
..} <- (Native -> Native) -> Par Native Native
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Native -> Native
forall t. t -> t
llvmTarget
Job
job <- Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJobUsing Seq (Int, sh, sh)
ranges Function
fun Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr ParamsR Native params
paramsR params
params Maybe Action
jobDone
Action -> Par Native ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Action -> Par Native ()) -> Action -> Par Native ()
forall a b. (a -> b) -> a -> b
$ Workers -> Job -> Action
schedule Workers
workers Job
job
{-# INLINABLE mkJob #-}
mkJob :: Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJob :: Int
-> Int
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJob Int
splits Int
minsize Function
fun Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr sh
from sh
to ParamsR Native params
paramsR params
params Maybe Action
jobDone =
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJobUsing (ShapeR sh
-> Int
-> Int
-> sh
-> sh
-> (Int -> sh -> sh -> (Int, sh, sh))
-> Seq (Int, sh, sh)
forall sh a.
ShapeR sh
-> Int -> Int -> sh -> sh -> (Int -> sh -> sh -> a) -> Seq a
divideWork ShapeR sh
shr Int
splits Int
minsize sh
from sh
to (,,)) Function
fun Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr ParamsR Native params
paramsR params
params Maybe Action
jobDone
{-# INLINABLE mkJobUsing #-}
mkJobUsing
:: Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJobUsing :: Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJobUsing Seq (Int, sh, sh)
ranges fun :: Function
fun@(ShortByteString
name,FunPtr ()
_) Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr ParamsR Native params
paramsR params
params Maybe Action
jobDone = do
Seq Action
jobTasks <- Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Par Native (Seq Action)
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Par Native (Seq Action)
mkTasksUsing Seq (Int, sh, sh)
ranges Function
fun Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr ParamsR Native params
paramsR params
params
IO Job -> Par Native Job
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Job -> Par Native Job) -> IO Job -> Par Native Job
forall a b. (a -> b) -> a -> b
$ ShortByteString -> Job -> IO Job
timed ShortByteString
name Job :: Seq Action -> Maybe Action -> Job
Job {Maybe Action
Seq Action
jobTasks :: Seq Action
jobDone :: Maybe Action
jobDone :: Maybe Action
jobTasks :: Seq Action
..}
{-# INLINABLE mkJobUsingIndex #-}
mkJobUsingIndex
:: Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJobUsingIndex :: Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Maybe Action
-> Par Native Job
mkJobUsingIndex Seq (Int, sh, sh)
ranges fun :: Function
fun@(ShortByteString
name,FunPtr ()
_) Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr ParamsR Native params
paramsR params
params Maybe Action
jobDone = do
Seq Action
jobTasks <- Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Par Native (Seq Action)
forall sh aenv params.
Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Par Native (Seq Action)
mkTasksUsingIndex Seq (Int, sh, sh)
ranges Function
fun Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr ParamsR Native params
paramsR params
params
IO Job -> Par Native Job
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Job -> Par Native Job) -> IO Job -> Par Native Job
forall a b. (a -> b) -> a -> b
$ ShortByteString -> Job -> IO Job
timed ShortByteString
name Job :: Seq Action -> Maybe Action -> Job
Job {Maybe Action
Seq Action
jobTasks :: Seq Action
jobDone :: Maybe Action
jobDone :: Maybe Action
jobTasks :: Seq Action
..}
{-# INLINABLE mkTasksUsing #-}
mkTasksUsing
:: Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Par Native (Seq Action)
mkTasksUsing :: Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Par Native (Seq Action)
mkTasksUsing Seq (Int, sh, sh)
ranges (ShortByteString
name, FunPtr ()
f) Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr ParamsR Native params
paramsR params
params = do
DList Arg
arg <- ParamsR Native (params, Val aenv)
-> (params, Val aenv) -> Par Native (DList (ArgR Native))
forall arch a.
Marshal arch =>
ParamsR arch a -> a -> Par arch (DList (ArgR arch))
marshalParams' @Native (ParamsR Native params
paramsR ParamsR Native params
-> TupR (ParamR Native) (Val aenv)
-> ParamsR Native (params, Val aenv)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Val aenv) -> TupR (ParamR Native) (Val aenv)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (Gamma aenv -> ParamR Native (Val aenv)
forall aenv arch. Gamma aenv -> ParamR arch (ValR arch aenv)
ParamRenv Gamma aenv
gamma)) (params
params, Val aenv
aenv)
Seq Action -> Par Native (Seq Action)
forall (m :: * -> *) a. Monad m => a -> m a
return (Seq Action -> Par Native (Seq Action))
-> Seq Action -> Par Native (Seq Action)
forall a b. (a -> b) -> a -> b
$ (((Int, sh, sh) -> Action) -> Seq (Int, sh, sh) -> Seq Action)
-> Seq (Int, sh, sh) -> ((Int, sh, sh) -> Action) -> Seq Action
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Int, sh, sh) -> Action) -> Seq (Int, sh, sh) -> Seq Action
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Seq (Int, sh, sh)
ranges (((Int, sh, sh) -> Action) -> Seq Action)
-> ((Int, sh, sh) -> Action) -> Seq Action
forall a b. (a -> b) -> a -> b
$ \(Int
_,sh
u,sh
v) -> do
String -> Action
sched (String -> Action) -> String -> Action
forall a b. (a -> b) -> a -> b
$ String -> String -> String -> String -> String
forall r. PrintfType r => String -> r
printf String
"%s (%s) -> (%s)" (ShortByteString -> String
S8.unpack ShortByteString
name) (ShapeR sh -> sh -> String
forall sh. ShapeR sh -> sh -> String
showShape ShapeR sh
shr sh
u) (ShapeR sh -> sh -> String
forall sh. ShapeR sh -> sh -> String
showShape ShapeR sh
shr sh
v)
let argU :: DList (ArgR Native)
argU = ShapeR sh -> sh -> DList (ArgR Native)
forall arch sh.
Marshal arch =>
ShapeR sh -> sh -> DList (ArgR arch)
marshalShape' @Native ShapeR sh
shr sh
u
let argV :: DList (ArgR Native)
argV = ShapeR sh -> sh -> DList (ArgR Native)
forall arch sh.
Marshal arch =>
ShapeR sh -> sh -> DList (ArgR arch)
marshalShape' @Native ShapeR sh
shr sh
v
FunPtr () -> RetType () -> [Arg] -> Action
forall a b. FunPtr a -> RetType b -> [Arg] -> IO b
callFFI FunPtr ()
f RetType ()
retVoid ([Arg] -> Action) -> [Arg] -> Action
forall a b. (a -> b) -> a -> b
$ DList Arg -> [Arg]
forall a. DList a -> [a]
DL.toList (DList Arg -> [Arg]) -> DList Arg -> [Arg]
forall a b. (a -> b) -> a -> b
$ DList (ArgR Native)
DList Arg
argU DList Arg -> DList Arg -> DList Arg
forall a. DList a -> DList a -> DList a
`DL.append` DList (ArgR Native)
DList Arg
argV DList Arg -> DList Arg -> DList Arg
forall a. DList a -> DList a -> DList a
`DL.append` DList Arg
arg
{-# INLINABLE mkTasksUsingIndex #-}
mkTasksUsingIndex
:: Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Par Native (Seq Action)
mkTasksUsingIndex :: Seq (Int, sh, sh)
-> Function
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> ParamsR Native params
-> params
-> Par Native (Seq Action)
mkTasksUsingIndex Seq (Int, sh, sh)
ranges (ShortByteString
name, FunPtr ()
f) Gamma aenv
gamma Val aenv
aenv ShapeR sh
shr ParamsR Native params
paramsR params
params = do
DList Arg
arg <- ParamsR Native (params, Val aenv)
-> (params, Val aenv) -> Par Native (DList (ArgR Native))
forall arch a.
Marshal arch =>
ParamsR arch a -> a -> Par arch (DList (ArgR arch))
marshalParams' @Native (ParamsR Native params
paramsR ParamsR Native params
-> TupR (ParamR Native) (Val aenv)
-> ParamsR Native (params, Val aenv)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ParamR Native (Val aenv) -> TupR (ParamR Native) (Val aenv)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (Gamma aenv -> ParamR Native (Val aenv)
forall aenv arch. Gamma aenv -> ParamR arch (ValR arch aenv)
ParamRenv Gamma aenv
gamma)) (params
params, Val aenv
aenv)
Seq Action -> Par Native (Seq Action)
forall (m :: * -> *) a. Monad m => a -> m a
return (Seq Action -> Par Native (Seq Action))
-> Seq Action -> Par Native (Seq Action)
forall a b. (a -> b) -> a -> b
$ (((Int, sh, sh) -> Action) -> Seq (Int, sh, sh) -> Seq Action)
-> Seq (Int, sh, sh) -> ((Int, sh, sh) -> Action) -> Seq Action
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Int, sh, sh) -> Action) -> Seq (Int, sh, sh) -> Seq Action
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Seq (Int, sh, sh)
ranges (((Int, sh, sh) -> Action) -> Seq Action)
-> ((Int, sh, sh) -> Action) -> Seq Action
forall a b. (a -> b) -> a -> b
$ \(Int
i,sh
u,sh
v) -> do
String -> Action
sched (String -> Action) -> String -> Action
forall a b. (a -> b) -> a -> b
$ String -> String -> String -> String -> String
forall r. PrintfType r => String -> r
printf String
"%s (%s) -> (%s)" (ShortByteString -> String
S8.unpack ShortByteString
name) (ShapeR sh -> sh -> String
forall sh. ShapeR sh -> sh -> String
showShape ShapeR sh
shr sh
u) (ShapeR sh -> sh -> String
forall sh. ShapeR sh -> sh -> String
showShape ShapeR sh
shr sh
v)
let argU :: DList (ArgR Native)
argU = ShapeR sh -> sh -> DList (ArgR Native)
forall arch sh.
Marshal arch =>
ShapeR sh -> sh -> DList (ArgR arch)
marshalShape' @Native ShapeR sh
shr sh
u
let argV :: DList (ArgR Native)
argV = ShapeR sh -> sh -> DList (ArgR Native)
forall arch sh.
Marshal arch =>
ShapeR sh -> sh -> DList (ArgR arch)
marshalShape' @Native ShapeR sh
shr sh
v
let argI :: DList Arg
argI = Arg -> DList Arg
forall a. a -> DList a
DL.singleton (Arg -> DList Arg) -> Arg -> DList Arg
forall a b. (a -> b) -> a -> b
$ Int -> ArgR Native
forall arch. Marshal arch => Int -> ArgR arch
marshalInt @Native Int
i
FunPtr () -> RetType () -> [Arg] -> Action
forall a b. FunPtr a -> RetType b -> [Arg] -> IO b
callFFI FunPtr ()
f RetType ()
retVoid ([Arg] -> Action) -> [Arg] -> Action
forall a b. (a -> b) -> a -> b
$ DList Arg -> [Arg]
forall a. DList a -> [a]
DL.toList (DList Arg -> [Arg]) -> DList Arg -> [Arg]
forall a b. (a -> b) -> a -> b
$ DList (ArgR Native)
DList Arg
argU DList Arg -> DList Arg -> DList Arg
forall a. DList a -> DList a -> DList a
`DL.append` DList (ArgR Native)
DList Arg
argV DList Arg -> DList Arg -> DList Arg
forall a. DList a -> DList a -> DList a
`DL.append` DList Arg
argI DList Arg -> DList Arg -> DList Arg
forall a. DList a -> DList a -> DList a
`DL.append` DList Arg
arg
memset :: Ptr Word8 -> Word8 -> Int -> IO ()
memset :: Ptr Word8 -> Word8 -> Int -> Action
memset Ptr Word8
p Word8
w Int
s = Ptr Word8 -> CInt -> CSize -> IO (Ptr Word8)
c_memset Ptr Word8
p (Word8 -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
w) (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
s) IO (Ptr Word8) -> Action -> Action
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> Action
forall (m :: * -> *) a. Monad m => a -> m a
return ()
foreign import ccall unsafe "string.h memset" c_memset
:: Ptr Word8 -> CInt -> CSize -> IO (Ptr Word8)
timed :: ShortByteString -> Job -> IO Job
timed :: ShortByteString -> Job -> IO Job
timed ShortByteString
name Job
job =
case HasInitialValue
Debug.debuggingIsEnabled of
HasInitialValue
False -> Job -> IO Job
forall (m :: * -> *) a. Monad m => a -> m a
return Job
job
HasInitialValue
True -> do
HasInitialValue
yes <- if HasInitialValue
Debug.monitoringIsEnabled
then HasInitialValue -> IO HasInitialValue
forall (m :: * -> *) a. Monad m => a -> m a
return HasInitialValue
True
else Flag -> IO HasInitialValue
Debug.getFlag Flag
Debug.dump_exec
if HasInitialValue
yes
then do
IORef Double
ref1 <- Double -> IO (IORef Double)
forall a. a -> IO (IORef a)
newIORef Double
0
IORef Integer
ref2 <- Integer -> IO (IORef Integer)
forall a. a -> IO (IORef a)
newIORef Integer
0
let start :: Action
start = do !Double
wall0 <- IO Double
getMonotonicTime
!Integer
cpu0 <- IO Integer
getCPUTime
IORef Double -> Double -> Action
forall a. IORef a -> a -> Action
writeIORef IORef Double
ref1 Double
wall0
IORef Integer -> Integer -> Action
forall a. IORef a -> a -> Action
writeIORef IORef Integer
ref2 Integer
cpu0
end :: Action
end = do !Integer
cpu1 <- IO Integer
getCPUTime
!Double
wall1 <- IO Double
getMonotonicTime
!Double
wall0 <- IORef Double -> IO Double
forall a. IORef a -> IO a
readIORef IORef Double
ref1
!Integer
cpu0 <- IORef Integer -> IO Integer
forall a. IORef a -> IO a
readIORef IORef Integer
ref2
let wallTime :: Double
wallTime = Double
wall1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
wall0
cpuTime :: Double
cpuTime = Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer
cpu1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
cpu0) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
1E-12
Processor -> Double -> Action
Debug.addProcessorTime Processor
Debug.Native Double
cpuTime
Flag -> String -> Action
Debug.traceIO Flag
Debug.dump_exec (String -> Action) -> String -> Action
forall a b. (a -> b) -> a -> b
$ String -> String -> String -> String
forall r. PrintfType r => String -> r
printf String
"exec: %s %s" (ShortByteString -> String
S8.unpack ShortByteString
name) (Double -> Double -> String
Debug.elapsedP Double
wallTime Double
cpuTime)
Job -> IO Job
forall (m :: * -> *) a. Monad m => a -> m a
return (Job -> IO Job) -> Job -> IO Job
forall a b. (a -> b) -> a -> b
$ Job :: Seq Action -> Maybe Action -> Job
Job { jobTasks :: Seq Action
jobTasks = Action
start Action -> Seq Action -> Seq Action
forall a. a -> Seq a -> Seq a
Seq.<| Job -> Seq Action
jobTasks Job
job
, jobDone :: Maybe Action
jobDone = case Job -> Maybe Action
jobDone Job
job of
Maybe Action
Nothing -> Action -> Maybe Action
forall a. a -> Maybe a
Just Action
end
Just Action
finished -> Action -> Maybe Action
forall a. a -> Maybe a
Just (Action
finished Action -> Action -> Action
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Action
end)
}
else
Job -> IO Job
forall (m :: * -> *) a. Monad m => a -> m a
return Job
job
foreign import ccall unsafe "clock_gettime_monotonic_seconds" getMonotonicTime :: IO Double
sched :: String -> IO ()
sched :: String -> Action
sched String
msg
= Flag -> Action -> Action
forall (m :: * -> *). MonadIO m => Flag -> m () -> m ()
Debug.when Flag
Debug.verbose
(Action -> Action) -> Action -> Action
forall a b. (a -> b) -> a -> b
$ Flag -> Action -> Action
forall (m :: * -> *). MonadIO m => Flag -> m () -> m ()
Debug.when Flag
Debug.dump_sched
(Action -> Action) -> Action -> Action
forall a b. (a -> b) -> a -> b
$ do ThreadId
tid <- IO ThreadId
myThreadId
String -> Action
Debug.putTraceMsg (String -> Action) -> String -> Action
forall a b. (a -> b) -> a -> b
$ String -> String -> String -> String
forall r. PrintfType r => String -> r
printf String
"sched: %s %s" (ThreadId -> String
forall a. Show a => a -> String
show ThreadId
tid) String
msg