{-# 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
-- Copyright   : [2014..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--

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) #-}

-- Array expression evaluation
-- ---------------------------

-- Computations are evaluated by traversing the AST bottom up, and for each node
-- distinguishing between three cases:
--
--  1. If it is a Use node, we return a reference to the array data. Even though
--     we execute with multiple cores, we assume a shared memory multiprocessor
--     machine.
--
--  2. If it is a non-skeleton node, such as a let binding or shape conversion,
--     then execute directly by updating the environment or similar.
--
--  3. If it is a skeleton node, then we need to execute the generated LLVM
--     code.
--
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


-- Skeleton implementation
-- -----------------------

-- Simple kernels just needs to know the shape of the output array.
--
{-# 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   -- XXX: must not unload the object code early
  Future (Array sh e) -> Par Native (Future (Array sh e))
forall (m :: * -> *) a. Monad m => a -> m a
return Future (Array sh e)
future

-- Mapping over an array can ignore the dimensionality of the array and
-- treat it as its underlying linear representation.
--
{-# 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)

-- Note: [Reductions]
--
-- There are two flavours of reduction:
--
--   1. If we are collapsing to a single value, then threads reduce strips of
--      the input in parallel, and then a single thread reduces the partial
--      reductions to a single value. Load balancing occurs over the input
--      stripes.
--
--   2. If this is a multidimensional reduction, then each inner dimension is
--      handled by a single thread. Load balancing occurs over the outer
--      dimension indices.
--
-- The entry points to executing the reduction are 'foldOp' and 'fold1Op', for
-- exclusive and inclusive reductions respectively. These functions handle
-- whether the input array is empty. If the input and output arrays are
-- non-empty, we then further dispatch (via 'foldCore') to 'foldAllOp' or
-- 'foldDimOp' for single or multidimensional reductions, respectively.
-- 'foldAllOp' in particular behaves differently whether we are evaluating the
-- array in parallel or sequentially.
--

{-# 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    -- empty, but possibly with non-zero dimensions
      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        -- output size of innermost dimension
    -> 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
    -- This is a multidimensional scan. Each partial scan result is evaluated
    -- individually by a thread, so no inter-thread communication is required.
    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

    -- This is a one-dimensional scan. If the array is small just compute it
    -- sequentially using a single thread, otherwise we require multiple steps
    -- to execute it in parallel.
    else
      if Int
n Int -> Int -> HasInitialValue
forall a. Ord a => a -> a -> HasInitialValue
< Int
8192
        -- sequential execution
        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

        -- parallel execution
        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
          --
          -- XXX: Should the sequential scan of the carry-in values just be
          -- executed immediately as part of the finalisation action?
          --
          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
    -- This is a multidimensional scan. Each partial scan result is evaluated
    -- individually by a thread, so no inter-thread communication is required.
    --
    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

    -- One dimensional scan. If the array is small just compute it sequentially
    -- with a single thread, otherwise we require multiple steps to execute it
    -- in parallel.
    --
    else
      if Int
n Int -> Int -> HasInitialValue
forall a. Ord a => a -> a -> HasInitialValue
< Int
8192
        -- sequential execution
        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

        -- parallel execution
        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

-- Forward permutation, specified by an indexing mapping into an array and a
-- combination function to combine elements.
--
{-# 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
    -- sequential execution does not require handling critical sections
    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

    -- parallel execution
    else
      case ShortByteString -> Lifetime FunctionTable -> Maybe Function
lookupFunction ShortByteString
"permuteP_rmw" Lifetime FunctionTable
nativeExecutable of
        -- using atomic operations
        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

        -- uses a temporary array of spin-locks to guard the critical section
        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                       -- border dimensions (i.e. index of first interior element)
    -> sh                       -- output array size
    -> 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

-- Compute the stencil border regions, where we may need to evaluate the
-- boundary conditions.
--
{-# 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


-- Skeleton execution
-- ------------------

(!#) :: 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

-- Schedule an operation over the entire iteration space, specifying the number
-- of partitions and minimum dimension size.
--
{-# INLINABLE scheduleOpWith #-}
scheduleOpWith
    :: Int            -- # subdivisions (hint)
    -> Int            -- minimum size of a dimension (must be a power of two)
    -> Function       -- function to execute
    -> Gamma aenv
    -> Val aenv
    -> ShapeR sh
    -> sh
    -> ParamsR Native params
    -> params
    -> Maybe Action   -- run after the last piece completes
    -> 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


-- Standard C functions
-- --------------------

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)


-- Debugging
-- ---------

-- Since the (new) thread scheduler does not operate in block-synchronous mode,
-- it is a bit more difficult to track how long an individual operation took to
-- execute as we won't know when exactly it will begin. The following method
-- (where initial timing information is recorded as the first task) should give
-- reasonable results.
--
-- TLM: missing GC stats information (verbose mode) since we aren't using the
--      the default 'timed' helper.
--
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

-- accelerate/cbits/clock.c
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