packages feed

accelerate-llvm-native-1.3.0.0: src/Data/Array/Accelerate/LLVM/Native/Execute.hs

{-# 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           = mapOp
  generate      = generateOp
  transform     = transformOp
  backpermute   = backpermuteOp
  fold True     = foldOp
  fold False    = fold1Op
  foldSeg i _   = foldSegOp i
  scan _ True   = scanOp
  scan _ False  = scan1Op
  scan' _       = scan'Op
  permute       = permuteOp
  stencil1      = stencil1Op
  stencil2      = stencil2Op
  aforeign      = 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 name repr NativeR{..} gamma aenv sh = do
  let fun   = nativeExecutable !# name
      param = TupRsingle $ ParamRarray repr
  Native{..} <- gets llvmTarget
  future     <- new
  result     <- allocateRemote repr sh
  scheduleOp fun gamma aenv (arrayRshape repr) sh param result
    `andThen` do putIO workers future result
                 touchLifetime nativeExecutable   -- XXX: must not unload the object code early
  return 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 inplace repr tp NativeR{..} gamma aenv input = do
  let fun   = nativeExecutable !# "map"
      sh    = shape input
      shr   = arrayRshape repr
      repr' = ArrayR shr tp
      param = TupRsingle (ParamRarray repr') `TupRpair` TupRsingle (ParamRarray repr)
  Native{..} <- gets llvmTarget
  future     <- new
  result     <- case inplace of
                  Just Refl -> return input
                  Nothing   -> allocateRemote repr' sh
  scheduleOp fun gamma aenv dim1 ((), size shr sh) param (result, input)
    `andThen` do putIO workers future result
                 touchLifetime nativeExecutable
  return future

{-# INLINE generateOp #-}
generateOp
    :: HasCallStack
    => ArrayR (Array sh e)
    -> ExecutableR Native
    -> Gamma aenv
    -> Val aenv
    -> sh
    -> Par Native (Future (Array sh e))
generateOp = simpleOp "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 repr repr' NativeR{..} gamma aenv sh' input = do
  let fun = nativeExecutable !# "transform"
  Native{..} <- gets llvmTarget
  future     <- new
  result     <- allocateRemote repr' sh'
  let param = TupRsingle (ParamRarray repr') `TupRpair` TupRsingle (ParamRarray repr)
  scheduleOp fun gamma aenv (arrayRshape repr') sh' param (result, input)
    `andThen` do putIO workers future result
                 touchLifetime nativeExecutable
  return 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 shr tp) shr' = transformOp (ArrayR shr tp) (ArrayR shr' 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 repr exe gamma aenv arr@(delayedShape -> sh@(sx, sz))
  = boundsCheck "empty array" (sz > 0)
  $ case size (ShapeRsnoc $ arrayRshape repr) sh of
      0 -> newFull =<< allocateRemote repr sx    -- empty, but possibly with non-zero dimensions
      _ -> foldCore repr exe gamma aenv 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 repr exe gamma aenv arr@(delayedShape -> sh@(sx, _)) =
  case size (ShapeRsnoc $ arrayRshape repr) sh of
    0 -> generateOp repr exe gamma aenv sx
    _ -> foldCore repr exe gamma aenv 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 repr exe gamma aenv arr
  | ArrayR ShapeRz tp <- repr
  = foldAllOp tp exe gamma aenv arr
  --
  | otherwise
  = foldDimOp repr exe gamma aenv arr

{-# INLINE foldAllOp #-}
foldAllOp
    :: HasCallStack
    => TypeR e
    -> ExecutableR Native
    -> Gamma aenv
    -> Val aenv
    -> Delayed (Vector e)
    -> Par Native (Future (Scalar e))
foldAllOp tp NativeR{..} gamma aenv arr = do
  Native{..}  <- gets llvmTarget
  future      <- new
  result      <- allocateRemote (ArrayR dim0 tp) ()
  let
      minsize = 4096
      splits  = numWorkers workers
      ranges  = divideWork1 splits minsize ((), 0) sh (,,)
      steps   = Seq.length ranges
      sh      = delayedShape arr
  --
  if steps <= 1
    then
      let param = TupRsingle (ParamRarray $ ArrayR dim0 tp) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray $ ArrayR dim1 tp)
      in  scheduleOpUsing ranges (nativeExecutable !# "foldAllS") gamma aenv dim1 param (result, manifest arr)
            `andThen` do putIO workers future result
                         touchLifetime nativeExecutable

    else do
      let param1 = TupRsingle (ParamRarray $ ArrayR dim1 tp) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray $ ArrayR dim1 tp)
      let param2 = TupRsingle (ParamRarray $ ArrayR dim1 tp) `TupRpair` TupRsingle (ParamRarray $ ArrayR dim0 tp)
      tmp   <- allocateRemote (ArrayR dim1 tp) ((), steps)
      job2  <- mkJobUsing (Seq.singleton (0, ((), 0), ((), steps))) (nativeExecutable !# "foldAllP2") gamma aenv dim1 param2 (tmp, result)
                 `andThen` do putIO workers future result
                              touchLifetime nativeExecutable

      job1  <- mkJobUsingIndex ranges (nativeExecutable !# "foldAllP1") gamma aenv dim1 param1 (tmp, manifest arr)
                 `andThen` do schedule workers job2

      liftIO $ schedule workers job1
  --
  return 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 repr NativeR{..} gamma aenv arr@(delayedShape -> (sh, _)) = do
  Native{..}  <- gets llvmTarget
  future      <- new
  result      <- allocateRemote repr sh
  let
      ArrayR shr tp = repr
      fun     = nativeExecutable !# "fold"
      splits  = numWorkers workers
      minsize = 1
      param   = TupRsingle (ParamRarray repr) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray $ ArrayR (ShapeRsnoc shr) tp)
  --
  scheduleOpWith splits minsize fun gamma aenv shr sh param (result, manifest arr)
    `andThen` do putIO workers future result
                 touchLifetime nativeExecutable
  return 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 int repr NativeR{..} gamma aenv input@(delayedShape -> (sh, _)) segments@(delayedShape -> ((), ss)) = do
  Native{..}  <- gets llvmTarget
  future      <- new
  let
      n       = ss-1
      splits  = numWorkers workers
      minsize = 1
      shR     = arrayRshape repr
      segR    = ArrayR dim1 $ TupRsingle $ SingleScalarType $ NumSingleType $ IntegralNumType int
      param   = TupRsingle (ParamRarray repr) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray segR)
  --
  result  <- allocateRemote repr (sh, n)
  scheduleOpWith splits minsize (nativeExecutable !# "foldSegP") gamma aenv shR (sh, n) param ((result, manifest input), manifest segments)
    `andThen` do putIO workers future result
                 touchLifetime nativeExecutable

  return 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 repr exe gamma aenv arr@(delayedShape -> (sz, n)) =
  case n of
    0 -> generateOp repr exe gamma aenv (sz, 1)
    _ -> scanCore   repr exe gamma aenv (n+1) 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 repr exe gamma aenv arr@(delayedShape -> (_, n))
  = boundsCheck "empty array" (n > 0)
  $ scanCore repr exe gamma aenv n 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 repr NativeR{..} gamma aenv m input@(delayedShape -> (sz, n)) = do
  Native{..}  <- gets llvmTarget
  future      <- new
  result      <- allocateRemote repr (sz, m)
  --
  let paramA = TupRsingle $ ParamRarray repr
      param  = paramA `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr)
      shR    = arrayRshape (reduceRank repr)

  if isMultiDim $ arrayRshape 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     = nativeExecutable !# "scanS"
          splits  = numWorkers workers
          minsize = 1
      in
      scheduleOpWith splits minsize fun gamma aenv shR sz param (result, manifest input)
        `andThen` do putIO workers future result
                     touchLifetime 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 n < 8192
        -- sequential execution
        then
          scheduleOpUsing (Seq.singleton (0, (), ())) (nativeExecutable !# "scanS") gamma aenv dim0 param (result, manifest input)
            `andThen` do putIO workers future result
                         touchLifetime nativeExecutable

        -- parallel execution
        else do
          let
              splits   = numWorkers workers
              minsize  = 8192
              ranges   = divideWork dim1 splits minsize ((), 0) ((), n) (,,)
              steps    = Seq.length ranges
              reprTmp  = ArrayR dim1 $ arrayRtype repr
              paramTmp = TupRsingle $ ParamRarray reprTmp
              param1   = TupRsingle ParamRint `TupRpair` paramA `TupRpair` paramTmp `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr)
              param3   = TupRsingle ParamRint `TupRpair` paramA `TupRpair` paramTmp
          --
          -- XXX: Should the sequential scan of the carry-in values just be
          -- executed immediately as part of the finalisation action?
          --
          tmp   <- allocateRemote (ArrayR dim1 $ arrayRtype repr) ((), steps)
          job3  <- mkJobUsingIndex ranges (nativeExecutable !# "scanP3") gamma aenv dim1 param3 ((steps, result), tmp)
                     `andThen` do putIO workers future result
                                  touchLifetime nativeExecutable
          job2  <- mkJobUsing (Seq.singleton (0, ((), 0), ((), steps))) (nativeExecutable !# "scanP2") gamma aenv dim1 paramTmp tmp
                     `andThen` schedule workers job3
          job1  <- mkJobUsingIndex ranges (nativeExecutable !# "scanP1") gamma aenv dim1 param1 (((steps, result), tmp), manifest input)
                     `andThen` schedule workers job2

          liftIO $ schedule workers job1
  --
  return 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 repr exe gamma aenv arr@(delayedShape -> (sz, n)) = do
  case n of
    0 -> do
      out     <- allocateRemote repr (sz, 0)
      sum     <- generateOp (reduceRank repr) exe gamma aenv sz
      future  <- new
      fork $ do sum' <- get sum
                put future (out, sum')
      return future
    --
    _ -> scan'Core repr exe gamma aenv 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 repr NativeR{..} gamma aenv input@(delayedShape -> sh@(sz, n)) = do
  let
      ArrayR shR eR   = repr
      ShapeRsnoc shR' = shR
      repr'           = ArrayR shR' eR
      paramA          = TupRsingle $ ParamRarray repr
      paramA'         = TupRsingle $ ParamRarray repr'
  --
  Native{..}  <- gets llvmTarget
  future      <- new
  result      <- allocateRemote repr  sh
  sums        <- allocateRemote repr' sz
  --
  if isMultiDim 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     = nativeExecutable !# "scanS"
          splits  = numWorkers workers
          minsize = 1
          param   = paramA `TupRpair` paramA' `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr)
      in
      scheduleOpWith splits minsize fun gamma aenv shR' sz param ((result, sums), manifest input)
        `andThen` do putIO workers future (result, sums)
                     touchLifetime 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 n < 8192
        -- sequential execution
        then
          let param = paramA `TupRpair` paramA' `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr)
          in  scheduleOpUsing (Seq.singleton (0, (), ())) (nativeExecutable !# "scanS") gamma aenv dim0 param ((result, sums), manifest input)
                `andThen` do putIO workers future (result, sums)
                             touchLifetime nativeExecutable

        -- parallel execution
        else do
          let
              splits   = numWorkers workers
              minsize  = 8192
              ranges   = divideWork1 splits minsize ((), 0) ((), n) (,,)
              steps    = Seq.length ranges
              reprTmp  = ArrayR dim1 eR
              paramTmp = TupRsingle $ ParamRarray reprTmp
              param1   = TupRsingle ParamRint `TupRpair` paramA `TupRpair` paramTmp `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr)
              param2   = paramA' `TupRpair` paramTmp
              param3   = TupRsingle ParamRint `TupRpair` paramA `TupRpair` paramTmp
          --
          tmp   <- allocateRemote reprTmp ((), steps)
          job3  <- mkJobUsingIndex ranges (nativeExecutable !# "scanP3") gamma aenv dim1 param3 ((steps, result), tmp)
                     `andThen` do putIO workers future (result, sums)
                                  touchLifetime nativeExecutable
          job2  <- mkJobUsing (Seq.singleton (0, ((), 0), ((), steps))) (nativeExecutable !# "scanP2") gamma aenv dim1 param2 (sums, tmp)
                     `andThen` schedule workers job3
          job1  <- mkJobUsingIndex ranges (nativeExecutable !# "scanP1") gamma aenv dim1 param1 (((steps, result), tmp), manifest input)
                     `andThen` schedule workers job2

          liftIO $ schedule workers job1
  --
  return future

isMultiDim :: ShapeR sh -> Bool
isMultiDim (ShapeRsnoc ShapeRz) = False
isMultiDim _                    = 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 inplace repr shr' NativeR{..} gamma aenv defaults@(shape -> shOut) input@(delayedShape -> shIn) = do
  let
      ArrayR shr tp = repr
      repr' = ArrayR shr' tp
  Native{..}  <- gets llvmTarget
  future      <- new
  result      <- if inplace
                   then Debug.trace Debug.dump_exec               "exec: permute/inplace"                            $ return defaults
                   else Debug.timed Debug.dump_exec (\wall cpu -> "exec: permute/clone " ++ Debug.elapsedS wall cpu) $ liftPar (cloneArray repr' defaults)
  let
      splits  = numWorkers workers
      minsize = case shr of
                  ShapeRsnoc ShapeRz              -> 4096
                  ShapeRsnoc (ShapeRsnoc ShapeRz) -> 64
                  _                               -> 16
      ranges  = divideWork shr splits minsize (empty shr) shIn (,,)
      steps   = Seq.length ranges
      paramR = TupRsingle (ParamRarray repr') `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr)
  --
  if steps <= 1
    -- sequential execution does not require handling critical sections
    then
      scheduleOpUsing ranges (nativeExecutable !# "permuteS") gamma aenv shr paramR (result, manifest input)
        `andThen` do putIO workers future result
                     touchLifetime nativeExecutable

    -- parallel execution
    else
      case lookupFunction "permuteP_rmw" nativeExecutable of
        -- using atomic operations
        Just f ->
          scheduleOpUsing ranges f gamma aenv shr paramR (result, manifest input)
            `andThen` do putIO workers future result
                         touchLifetime nativeExecutable

        -- uses a temporary array of spin-locks to guard the critical section
        Nothing -> do
          let m           = size shr' shOut
              reprBarrier = ArrayR dim1 $ TupRsingle scalarTypeWord8
              paramR'     = TupRsingle (ParamRarray repr') `TupRpair` TupRsingle (ParamRarray reprBarrier) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr)
          --
          barrier@(Array _ adb) <- allocateRemote reprBarrier ((), m) :: Par Native (Vector Word8)
          liftIO $ memset (unsafeUniqueArrayPtr adb) 0 m
          scheduleOpUsing ranges (nativeExecutable !# "permuteP_mutex") gamma aenv shr paramR' ((result, barrier), manifest input)
            `andThen` do putIO workers future result
                         touchLifetime nativeExecutable
  --
  return 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 tp repr halo exe gamma aenv input@(delayedShape -> sh) =
  stencilCore repr exe gamma aenv halo sh (TupRsingle $ ParamRmaybe $ ParamRarray $ ArrayR (arrayRshape repr) tp) (manifest 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 t1 t2 repr halo exe gamma aenv input1@(delayedShape -> sh1) input2@(delayedShape -> sh2) =
  stencilCore repr exe gamma aenv halo (intersect (arrayRshape repr) sh1 sh2) (param t1 `TupRpair` param t2) (manifest input1, manifest input2)
  where
    shr = arrayRshape repr
    param :: TypeR t -> ParamsR Native (Maybe (Array sh t))
    param = TupRsingle . ParamRmaybe . ParamRarray . ArrayR 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 repr NativeR{..} gamma aenv halo sh paramsR params = do
  Native{..} <- gets llvmTarget
  future     <- new
  result     <- allocateRemote repr sh
  let
      shr     = arrayRshape repr
      inside  = nativeExecutable !# "stencil_inside"
      border  = nativeExecutable !# "stencil_border"

      splits  = numWorkers workers
      minsize = case shr of
                  ShapeRsnoc ShapeRz              -> 4096
                  ShapeRsnoc (ShapeRsnoc ShapeRz) -> 64
                  _                               -> 16

      ins     = divideWork shr splits minsize halo (sub sh halo) (,,)
      outs    = asum . flip fmap (stencilBorders shr sh halo) $ \(u,v) -> divideWork shr splits minsize u v (,,)

      sub :: sh -> sh -> sh
      sub a b = go shr a b
        where
          go :: ShapeR t -> t -> t -> t
          go ShapeRz          ()      ()      = ()
          go (ShapeRsnoc shr') (xa,xb) (ya,yb) = (go shr' xa ya, xb - yb)

      paramsR' = TupRsingle (ParamRarray repr) `TupRpair` paramsR
  --
  jobsInside <- mkTasksUsing ins  inside gamma aenv shr paramsR' (result, params)
  jobsBorder <- mkTasksUsing outs border gamma aenv shr paramsR' (result, params)
  let jobTasks  = jobsInside Seq.>< jobsBorder
      jobDone   = Just $ do putIO workers future result
                            touchLifetime nativeExecutable
  --
  liftIO $ schedule workers =<< timed "stencil" Job{..}
  return 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 shr sh halo = Seq.fromFunction (2 * rank shr) face
  where
    face :: Int -> (sh, sh)
    face n = go n shr sh halo

    go :: Int -> ShapeR t -> t -> t -> (t, t)
    go _ ShapeRz          ()         ()         = ((), ())
    go n (ShapeRsnoc shr') (sha, sza) (shb, szb)
      = let
            (sha', shb')  = go (n-2) shr' sha shb
            (sza', szb')
              | n <  0    = (0,       sza)
              | n == 0    = (0,       szb)
              | n == 1    = (sza-szb, sza)
              | otherwise = (szb,     sza-szb)
        in
        ((sha', sza'), (shb', szb'))

{-# INLINE aforeignOp #-}
aforeignOp
    :: HasCallStack
    => String
    -> ArraysR as
    -> ArraysR bs
    -> (as -> Par Native (Future bs))
    -> as
    -> Par Native (Future bs)
aforeignOp name _ _ asm arr = do
  wallBegin <- liftIO getMonotonicTime
  result    <- Debug.timed Debug.dump_exec (\wall cpu -> printf "exec: %s %s" name (Debug.elapsedP wall cpu)) (asm arr)
  wallEnd   <- liftIO getMonotonicTime
  liftIO $ Debug.addProcessorTime Debug.Native (wallEnd - wallBegin)
  return result


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

(!#) :: HasCallStack => Lifetime FunctionTable -> ShortByteString -> Function
(!#) exe name
  = fromMaybe (internalError ("function not found: " ++ S8.unpack name))
  $ lookupFunction name exe

lookupFunction :: ShortByteString -> Lifetime FunctionTable -> Maybe Function
lookupFunction name nativeExecutable = do
  find (\(n,_) -> n == name) (functionTable (unsafeGetValue nativeExecutable))

andThen :: (Maybe a -> t) -> a -> t
andThen f g = f (Just g)

delayedShape :: Delayed (Array sh e) -> sh
delayedShape (Delayed sh) = sh
delayedShape (Manifest a) = shape a

manifest :: Delayed (Array sh e) -> Maybe (Array sh e)
manifest (Manifest a) = Just a
manifest Delayed{}    = Nothing


{-# INLINABLE scheduleOp #-}
scheduleOp
    :: HasCallStack
    => Function
    -> Gamma aenv
    -> Val aenv
    -> ShapeR sh
    -> sh
    -> ParamsR Native params
    -> params
    -> Maybe Action
    -> Par Native ()
scheduleOp fun gamma aenv shr sz paramsR params done = do
  Native{..} <- gets llvmTarget
  let
      splits  = numWorkers workers
      minsize = case shr of
                  ShapeRsnoc ShapeRz              -> 4096
                  ShapeRsnoc (ShapeRsnoc ShapeRz) -> 64
                  _                               -> 16
  --
  scheduleOpWith splits minsize fun gamma aenv shr sz paramsR params 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 splits minsize fun gamma aenv shr sz paramsR params done = do
  Native{..} <- gets llvmTarget
  job        <- mkJob splits minsize fun gamma aenv shr (empty shr) sz paramsR params done
  liftIO $ schedule workers job

{-# INLINABLE scheduleOpUsing #-}
scheduleOpUsing
    :: Seq (Int, sh, sh)
    -> Function
    -> Gamma aenv
    -> Val aenv
    -> ShapeR sh
    -> ParamsR Native params
    -> params
    -> Maybe Action
    -> Par Native ()
scheduleOpUsing ranges fun gamma aenv shr paramsR params jobDone = do
  Native{..} <- gets llvmTarget
  job        <- mkJobUsing ranges fun gamma aenv shr paramsR params jobDone
  liftIO $ schedule workers 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 splits minsize fun gamma aenv shr from to paramsR params jobDone =
  mkJobUsing (divideWork shr splits minsize from to (,,)) fun gamma aenv shr paramsR params 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 ranges fun@(name,_) gamma aenv shr paramsR params jobDone = do
  jobTasks <- mkTasksUsing ranges fun gamma aenv shr paramsR params
  liftIO    $ timed name Job {..}

{-# INLINABLE mkJobUsingIndex #-}
mkJobUsingIndex
      :: Seq (Int, sh, sh)
      -> Function
      -> Gamma aenv
      -> Val aenv
      -> ShapeR sh
      -> ParamsR Native params
      -> params
      -> Maybe Action
      -> Par Native Job
mkJobUsingIndex ranges fun@(name,_) gamma aenv shr paramsR params jobDone = do
  jobTasks <- mkTasksUsingIndex ranges fun gamma aenv shr paramsR params
  liftIO    $ timed name Job {..}

{-# INLINABLE mkTasksUsing #-}
mkTasksUsing
      :: Seq (Int, sh, sh)
      -> Function
      -> Gamma aenv
      -> Val aenv
      -> ShapeR sh
      -> ParamsR Native params
      -> params
      -> Par Native (Seq Action)
mkTasksUsing ranges (name, f) gamma aenv shr paramsR params = do
  arg <- marshalParams' @Native (paramsR `TupRpair` TupRsingle (ParamRenv gamma)) (params, aenv)
  return $ flip fmap ranges $ \(_,u,v) -> do
    sched $ printf "%s (%s) -> (%s)" (S8.unpack name) (showShape shr u) (showShape shr v)
    let argU = marshalShape' @Native shr u
    let argV = marshalShape' @Native shr v
    callFFI f retVoid $ DL.toList $ argU `DL.append` argV `DL.append` arg

{-# INLINABLE mkTasksUsingIndex #-}
mkTasksUsingIndex
      :: Seq (Int, sh, sh)
      -> Function
      -> Gamma aenv
      -> Val aenv
      -> ShapeR sh
      -> ParamsR Native params
      -> params
      -> Par Native (Seq Action)
mkTasksUsingIndex ranges (name, f) gamma aenv shr paramsR params = do
  arg <- marshalParams' @Native (paramsR `TupRpair` TupRsingle (ParamRenv gamma)) (params, aenv)
  return $ flip fmap ranges $ \(i,u,v) -> do
    sched $ printf "%s (%s) -> (%s)" (S8.unpack name) (showShape shr u) (showShape shr v)
    let argU = marshalShape' @Native shr u
    let argV = marshalShape' @Native shr v
    let argI = DL.singleton $ marshalInt @Native i
    callFFI f retVoid $ DL.toList $ argU `DL.append` argV `DL.append` argI `DL.append` arg


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

memset :: Ptr Word8 -> Word8 -> Int -> IO ()
memset p w s = c_memset p (fromIntegral w) (fromIntegral s) >> 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 name job =
  case Debug.debuggingIsEnabled of
    False -> return job
    True  -> do
      yes <- if Debug.monitoringIsEnabled
               then return True
               else Debug.getFlag Debug.dump_exec
      --
      if yes
        then do
          ref1 <- newIORef 0
          ref2 <- newIORef 0
          let start = do !wall0 <- getMonotonicTime
                         !cpu0  <- getCPUTime
                         writeIORef ref1 wall0
                         writeIORef ref2 cpu0

              end   = do !cpu1  <- getCPUTime
                         !wall1 <- getMonotonicTime
                         !wall0 <- readIORef ref1
                         !cpu0  <- readIORef ref2
                         --
                         let wallTime = wall1 - wall0
                             cpuTime  = fromIntegral (cpu1 - cpu0) * 1E-12
                         --
                         Debug.addProcessorTime Debug.Native cpuTime
                         Debug.traceIO Debug.dump_exec $ printf "exec: %s %s" (S8.unpack name) (Debug.elapsedP wallTime cpuTime)
              --
          return $ Job { jobTasks = start Seq.<| jobTasks job
                       , jobDone  = case jobDone job of
                                      Nothing       -> Just end
                                      Just finished -> Just (finished >> end)
                       }
        else
          return job

-- accelerate/cbits/clock.c
foreign import ccall unsafe "clock_gettime_monotonic_seconds" getMonotonicTime :: IO Double


sched :: String -> IO ()
sched msg
  = Debug.when Debug.verbose
  $ Debug.when Debug.dump_sched
  $ do tid <- myThreadId
       Debug.putTraceMsg $ printf "sched: %s %s" (show tid) msg