accelerate-llvm-ptx-1.3.0.0: src/Data/Array/Accelerate/LLVM/PTX/Execute.hs
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- |
-- Module : Data.Array.Accelerate.LLVM.PTX.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.PTX.Execute (
executeAcc,
executeOpenAcc,
) where
import Data.Array.Accelerate.Analysis.Match
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.PTX.Analysis.Launch ( multipleOf )
import Data.Array.Accelerate.LLVM.PTX.Array.Data
import Data.Array.Accelerate.LLVM.PTX.Array.Prim ( memsetArrayAsync )
import Data.Array.Accelerate.LLVM.PTX.Execute.Async
import Data.Array.Accelerate.LLVM.PTX.Execute.Environment
import Data.Array.Accelerate.LLVM.PTX.Execute.Marshal
import Data.Array.Accelerate.LLVM.PTX.Execute.Stream ( Stream )
import Data.Array.Accelerate.LLVM.PTX.Link
import Data.Array.Accelerate.LLVM.PTX.Target
import qualified Data.Array.Accelerate.LLVM.PTX.Debug as Debug
import qualified Data.Array.Accelerate.LLVM.PTX.Execute.Event as Event
import qualified Foreign.CUDA.Driver as CUDA
import Control.Monad ( when, forM_ )
import Control.Monad.Reader ( asks, local )
import Control.Monad.State ( liftIO )
import Data.ByteString.Short.Char8 ( ShortByteString, unpack )
import qualified Data.DList as DL
import Data.List ( find )
import Data.Maybe ( fromMaybe )
import Text.Printf ( printf )
import Prelude hiding ( exp, map, sum, scanl, scanr )
{-# SPECIALISE INLINE executeAcc :: ExecAcc PTX a -> Par PTX (FutureArraysR PTX a) #-}
{-# SPECIALISE INLINE executeOpenAcc :: ExecOpenAcc PTX aenv a -> Val aenv -> Par PTX (FutureArraysR PTX 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. The data
-- will already have been copied to the device during compilation of the
-- kernels.
--
-- 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 PTX 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 need to know the shape of the output array
--
{-# INLINE simpleOp #-}
simpleOp
:: HasCallStack
=> ShortByteString
-> ArrayR (Array sh e)
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> sh
-> Par PTX (Future (Array sh e))
simpleOp name repr exe gamma aenv sh =
withExecutable exe $ \ptxExecutable -> do
future <- new
result <- allocateRemote repr sh
--
let paramR = TupRsingle $ ParamRarray repr
executeOp (ptxExecutable !# name) gamma aenv (arrayRshape repr) sh paramR result
put future result
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 PTX
-> Gamma aenv
-> Val aenv
-> Array sh a
-> Par PTX (Future (Array sh b))
mapOp inplace repr tp exe gamma aenv input@(shape -> sh) =
withExecutable exe $ \ptxExecutable -> do
let reprOut = ArrayR (arrayRshape repr) tp
future <- new
result <- case inplace of
Just Refl -> return input
Nothing -> allocateRemote reprOut sh
--
let paramsR = TupRsingle (ParamRarray reprOut) `TupRpair` TupRsingle (ParamRarray repr)
executeOp (ptxExecutable !# "map") gamma aenv (arrayRshape repr) sh paramsR (result, input)
put future result
return future
{-# INLINE generateOp #-}
generateOp
:: HasCallStack
=> ArrayR (Array sh e)
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> sh
-> Par PTX (Future (Array sh e))
generateOp = simpleOp "generate"
{-# INLINE transformOp #-}
transformOp
:: HasCallStack
=> ArrayR (Array sh a)
-> ArrayR (Array sh' b)
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> sh'
-> Array sh a
-> Par PTX (Future (Array sh' b))
transformOp repr repr' exe gamma aenv sh' input =
withExecutable exe $ \ptxExecutable -> do
future <- new
result <- allocateRemote repr' sh'
let paramsR = TupRsingle (ParamRarray repr') `TupRpair` TupRsingle (ParamRarray repr)
executeOp (ptxExecutable !# "transform") gamma aenv (arrayRshape repr') sh' paramsR (result, input)
put future result
return future
{-# INLINE backpermuteOp #-}
backpermuteOp
:: HasCallStack
=> ArrayR (Array sh e)
-> ShapeR sh'
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> sh'
-> Array sh e
-> Par PTX (Future (Array sh' e))
backpermuteOp (ArrayR shr tp) shr' = transformOp (ArrayR shr tp) (ArrayR shr' tp)
-- There are two flavours of fold operation:
--
-- 1. If we are collapsing to a single value, then multiple thread blocks are
-- working together. Since thread blocks synchronise with each other via
-- kernel launches, each block computes a partial sum and the kernel is
-- launched recursively until the final value is reached.
--
-- 2. If this is a multidimensional reduction, then each inner dimension is
-- handled by a single thread block, so no global communication is
-- necessary. Furthermore are two kernel flavours: each innermost dimension
-- can be cooperatively reduced by (a) a thread warp; or (b) a thread
-- block. Currently we always use the first, but require benchmarking to
-- determine when to select each.
--
{-# INLINE fold1Op #-}
fold1Op
:: HasCallStack
=> ArrayR (Array sh e)
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par PTX (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 one or more non-zero dimensions
_ -> foldCore repr exe gamma aenv arr
{-# INLINE foldOp #-}
foldOp
:: HasCallStack
=> ArrayR (Array sh e)
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par PTX (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 PTX
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par PTX (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
:: forall aenv e. HasCallStack
=> TypeR e
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Delayed (Vector e)
-> Par PTX (Future (Scalar e))
foldAllOp tp exe gamma aenv input =
withExecutable exe $ \ptxExecutable -> do
future <- new
let
ks = ptxExecutable !# "foldAllS"
km1 = ptxExecutable !# "foldAllM1"
km2 = ptxExecutable !# "foldAllM2"
sh@((), n) = delayedShape input
paramsRinput = TupRsingle $ ParamRmaybe $ ParamRarray $ ArrayR dim1 tp
paramsRdim0 = TupRsingle $ ParamRarray $ ArrayR dim0 tp
paramsRdim1 = TupRsingle $ ParamRarray $ ArrayR dim1 tp
--
if kernelThreadBlocks ks n == 1
then do
-- The array is small enough that we can compute it in a single step
result <- allocateRemote (ArrayR dim0 tp) ()
let paramsR = paramsRdim0 `TupRpair` paramsRinput
executeOp ks gamma aenv dim1 sh paramsR (result, manifest input)
put future result
else do
-- Multi-kernel reduction to a single element. The first kernel integrates
-- any delayed elements, and the second is called recursively until
-- reaching a single element.
let
rec :: Vector e -> Par PTX ()
rec tmp@(Array ((),m) adata)
| m <= 1 = put future (Array () adata)
| otherwise = do
let sh' = ((), m `multipleOf` kernelThreadBlockSize km2)
out <- allocateRemote (ArrayR dim1 tp) sh'
let paramsR2 = paramsRdim1 `TupRpair` paramsRdim1
executeOp km2 gamma aenv dim1 sh' paramsR2 (tmp, out)
rec out
--
let sh' = ((), n `multipleOf` kernelThreadBlockSize km1)
tmp <- allocateRemote (ArrayR dim1 tp) sh'
let paramsR1 = paramsRdim1 `TupRpair` paramsRinput
executeOp km1 gamma aenv dim1 sh' paramsR1 (tmp, manifest input)
rec tmp
--
return future
{-# INLINE foldDimOp #-}
foldDimOp
:: HasCallStack
=> ArrayR (Array sh e)
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par PTX (Future (Array sh e))
foldDimOp repr@(ArrayR shr tp) exe gamma aenv input@(delayedShape -> (sh, sz))
| sz == 0 = generateOp repr exe gamma aenv sh
| otherwise =
withExecutable exe $ \ptxExecutable -> do
future <- new
result <- allocateRemote repr sh
--
let paramsR = TupRsingle (ParamRarray repr) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray $ ArrayR (ShapeRsnoc shr) tp)
executeOp (ptxExecutable !# "fold") gamma aenv shr sh paramsR (result, manifest input)
put future result
return future
{-# INLINE foldSegOp #-}
foldSegOp
:: HasCallStack
=> IntegralType i
-> ArrayR (Array (sh, Int) e)
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Delayed (Segments i)
-> Par PTX (Future (Array (sh, Int) e))
foldSegOp intTp repr exe gamma aenv input@(delayedShape -> (sh, sz)) segments@(delayedShape -> ((), ss)) =
withExecutable exe $ \ptxExecutable -> do
let
ArrayR (ShapeRsnoc shr') _ = repr
reprSeg = ArrayR dim1 $ TupRsingle $ SingleScalarType $ NumSingleType $ IntegralNumType intTp
n = ss - 1 -- segments array has been 'scanl (+) 0'`ed
m = size shr' sh * n
foldseg = if (sz`quot`ss) < (2 * kernelThreadBlockSize foldseg_cta)
then foldseg_warp
else foldseg_cta
--
foldseg_cta = ptxExecutable !# "foldSeg_block"
foldseg_warp = ptxExecutable !# "foldSeg_warp"
-- qinit = ptxExecutable !# "qinit"
--
future <- new
result <- allocateRemote repr (sh, n)
let paramsR = TupRsingle (ParamRarray repr) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray reprSeg)
executeOp foldseg gamma aenv dim1 ((), m) paramsR ((result, manifest input), manifest segments)
put future result
return future
{-# INLINE scanOp #-}
scanOp
:: HasCallStack
=> ArrayR (Array (sh, Int) e)
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par PTX (Future (Array (sh, Int) e))
scanOp repr exe gamma aenv input@(delayedShape -> (sz, n)) =
case n of
0 -> generateOp repr exe gamma aenv (sz, 1)
_ -> scanCore repr exe gamma aenv (n+1) input
{-# INLINE scan1Op #-}
scan1Op
:: HasCallStack
=> ArrayR (Array (sh, Int) e)
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par PTX (Future (Array (sh, Int) e))
scan1Op repr exe gamma aenv input@(delayedShape -> (_, n))
= boundsCheck "empty array" (n > 0)
$ scanCore repr exe gamma aenv n input
{-# INLINE scanCore #-}
scanCore
:: HasCallStack
=> ArrayR (Array (sh, Int) e)
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Int -- output size of innermost dimension
-> Delayed (Array (sh, Int) e)
-> Par PTX (Future (Array (sh, Int) e))
scanCore repr exe gamma aenv m input
| ArrayR (ShapeRsnoc ShapeRz) tp <- repr
= scanAllOp tp exe gamma aenv m input
--
| otherwise
= scanDimOp repr exe gamma aenv m input
{-# INLINE scanAllOp #-}
scanAllOp
:: HasCallStack
=> TypeR e
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Int -- output size
-> Delayed (Vector e)
-> Par PTX (Future (Vector e))
scanAllOp tp exe gamma aenv m input@(delayedShape -> ((), n)) =
withExecutable exe $ \ptxExecutable -> do
let
k1 = ptxExecutable !# "scanP1"
k2 = ptxExecutable !# "scanP2"
k3 = ptxExecutable !# "scanP3"
--
c = kernelThreadBlockSize k1
s = n `multipleOf` c
--
repr = ArrayR dim1 tp
paramR = TupRsingle $ ParamRarray repr
paramsR1 = paramR `TupRpair` paramR `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr)
paramsR3 = paramR `TupRpair` paramR `TupRpair` TupRsingle ParamRint
--
future <- new
result <- allocateRemote repr ((), m)
-- Step 1: Independent thread-block-wide scans of the input. Small arrays
-- which can be computed by a single thread block will require no
-- additional work.
tmp <- allocateRemote repr ((), s)
executeOp k1 gamma aenv dim1 ((), s) paramsR1 ((tmp, result), manifest input)
-- Step 2: Multi-block reductions need to compute the per-block prefix,
-- then apply those values to the partial results.
when (s > 1) $ do
executeOp k2 gamma aenv dim1 ((), s) paramR tmp
executeOp k3 gamma aenv dim1 ((), s-1) paramsR3 ((tmp, result), c)
put future result
return future
{-# INLINE scanDimOp #-}
scanDimOp
:: HasCallStack
=> ArrayR (Array (sh, Int) e)
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Int
-> Delayed (Array (sh, Int) e)
-> Par PTX (Future (Array (sh, Int) e))
scanDimOp repr exe gamma aenv m input@(delayedShape -> (sz, _)) =
withExecutable exe $ \ptxExecutable -> do
let ArrayR (ShapeRsnoc shr') _ = repr
future <- new
result <- allocateRemote repr (sz, m)
let paramsR = TupRsingle (ParamRarray repr) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr)
executeOp (ptxExecutable !# "scan") gamma aenv dim1 ((), size shr' sz) paramsR (result, manifest input)
put future result
return future
{-# INLINE scan'Op #-}
scan'Op
:: HasCallStack
=> ArrayR (Array (sh, Int) e)
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par PTX (Future (Array (sh, Int) e, Array sh e))
scan'Op repr exe gamma aenv input@(delayedShape -> (sz, n)) =
case n of
0 -> do
future <- new
result <- allocateRemote repr (sz, 0)
sums <- generateOp (reduceRank repr) exe gamma aenv sz
fork $ do sums' <- get sums
put future (result, sums')
return future
--
_ -> scan'Core repr exe gamma aenv input
{-# INLINE scan'Core #-}
scan'Core
:: HasCallStack
=> ArrayR (Array (sh, Int) e)
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par PTX (Future (Array (sh, Int) e, Array sh e))
scan'Core repr exe gamma aenv input
| ArrayR (ShapeRsnoc ShapeRz) tp <- repr
= scan'AllOp tp exe gamma aenv input
--
| otherwise
= scan'DimOp repr exe gamma aenv input
{-# INLINE scan'AllOp #-}
scan'AllOp
:: HasCallStack
=> TypeR e
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Delayed (Vector e)
-> Par PTX (Future (Vector e, Scalar e))
scan'AllOp tp exe gamma aenv input@(delayedShape -> ((), n)) =
withExecutable exe $ \ptxExecutable -> do
let
repr = ArrayR dim1 tp
paramRdim0 = TupRsingle $ ParamRarray $ ArrayR dim0 tp
paramRdim1 = TupRsingle $ ParamRarray repr
k1 = ptxExecutable !# "scanP1"
k2 = ptxExecutable !# "scanP2"
k3 = ptxExecutable !# "scanP3"
--
c = kernelThreadBlockSize k1
s = n `multipleOf` c
--
future <- new
result <- allocateRemote repr ((), n)
tmp <- allocateRemote repr ((), s)
-- Step 1: independent thread-block-wide scans. Each block stores its partial
-- sum to a temporary array.
let paramsR1 = paramRdim1 `TupRpair` paramRdim1 `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr)
executeOp k1 gamma aenv dim1 ((), s) paramsR1 ((tmp, result), manifest input)
-- If this was a small array that was processed by a single thread block then
-- we are done, otherwise compute the per-block prefix and apply those values
-- to the partial results.
if s == 1
then
case tmp of
Array _ ad -> put future (result, Array () ad)
else do
sums <- allocateRemote (ArrayR dim0 tp) ()
let paramsR2 = paramRdim1 `TupRpair` paramRdim0
let paramsR3 = paramRdim1 `TupRpair` paramRdim1 `TupRpair` TupRsingle ParamRint
executeOp k2 gamma aenv dim1 ((), s) paramsR2 (tmp, sums)
executeOp k3 gamma aenv dim1 ((), s-1) paramsR3 ((tmp, result), c)
put future (result, sums)
--
return future
{-# INLINE scan'DimOp #-}
scan'DimOp
:: HasCallStack
=> ArrayR (Array (sh, Int) e)
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Delayed (Array (sh, Int) e)
-> Par PTX (Future (Array (sh, Int) e, Array sh e))
scan'DimOp repr@(ArrayR (ShapeRsnoc shr') _) exe gamma aenv input@(delayedShape -> sh@(sz, _)) =
withExecutable exe $ \ptxExecutable -> do
future <- new
result <- allocateRemote repr sh
sums <- allocateRemote (reduceRank repr) sz
let paramsR = TupRsingle (ParamRarray repr) `TupRpair` TupRsingle (ParamRarray $ reduceRank repr) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr)
executeOp (ptxExecutable !# "scan") gamma aenv dim1 ((), size shr' sz) paramsR ((result, sums), manifest input)
put future (result, sums)
return future
{-# INLINE permuteOp #-}
permuteOp
:: HasCallStack
=> Bool
-> ArrayR (Array sh e)
-> ShapeR sh'
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Array sh' e
-> Delayed (Array sh e)
-> Par PTX (Future (Array sh' e))
permuteOp inplace repr@(ArrayR shr tp) shr' exe gamma aenv defaults@(shape -> shOut) input@(delayedShape -> shIn) =
withExecutable exe $ \ptxExecutable -> do
let
n = size shr shIn
m = size shr' shOut
repr' = ArrayR shr' tp
reprLock = ArrayR dim1 $ TupRsingle $ scalarTypeWord32
paramR = TupRsingle $ ParamRmaybe $ ParamRarray repr
paramR' = TupRsingle $ ParamRarray repr'
kernel = case functionTable ptxExecutable of
k:_ -> k
_ -> internalError "no kernels found"
--
future <- new
result <- if inplace
then Debug.trace Debug.dump_exec "exec: permute/inplace" $ return defaults
else Debug.trace Debug.dump_exec "exec: permute/clone" $ get =<< cloneArrayAsync repr' defaults
--
case kernelName kernel of
-- execute directly using atomic operations
"permute_rmw" ->
let paramsR = paramR' `TupRpair` paramR
in executeOp kernel gamma aenv dim1 ((), n) paramsR (result, manifest input)
-- a temporary array is required for spin-locks around the critical section
"permute_mutex" -> do
barrier <- new :: Par PTX (Future (Vector Word32))
Array _ ad <- allocateRemote reprLock ((), m)
fork $ do fill <- memsetArrayAsync (NumSingleType $ IntegralNumType TypeWord32) m 0 ad
put barrier . Array ((), m) =<< get fill
--
let paramsR = paramR' `TupRpair` TupRsingle (ParamRfuture $ ParamRarray reprLock) `TupRpair` paramR
executeOp kernel gamma aenv dim1 ((), n) paramsR ((result, barrier), manifest input)
_ -> internalError "unexpected kernel image"
--
put future result
return future
{-# INLINE stencil1Op #-}
stencil1Op
:: HasCallStack
=> TypeR a
-> ArrayR (Array sh b)
-> sh
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Delayed (Array sh a)
-> Par PTX (Future (Array sh b))
stencil1Op tp repr@(ArrayR shr _) halo exe gamma aenv input@(delayedShape -> sh) =
stencilCore repr exe gamma aenv halo sh paramsR (manifest input)
where paramsR = TupRsingle $ ParamRmaybe $ ParamRarray $ ArrayR shr tp
-- Using the defaulting instances for stencil operations (for now).
--
{-# INLINE stencil2Op #-}
stencil2Op
:: HasCallStack
=> TypeR a
-> TypeR b
-> ArrayR (Array sh c)
-> sh
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> Delayed (Array sh a)
-> Delayed (Array sh b)
-> Par PTX (Future (Array sh c))
stencil2Op tpA tpB repr@(ArrayR shr _) halo exe gamma aenv input1@(delayedShape -> sh1) input2@(delayedShape -> sh2) =
stencilCore repr exe gamma aenv halo (intersect (arrayRshape repr) sh1 sh2) paramsR (manifest input1, manifest input2)
where paramsR = TupRsingle (ParamRmaybe $ ParamRarray $ ArrayR shr tpA) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray $ ArrayR shr tpB)
{-# INLINE stencilCore #-}
stencilCore
:: forall aenv sh e params. HasCallStack
=> ArrayR (Array sh e)
-> ExecutableR PTX
-> Gamma aenv
-> Val aenv
-> sh -- border dimensions (i.e. index of first interior element)
-> sh -- output array size
-> ParamsR PTX params
-> params
-> Par PTX (Future (Array sh e))
stencilCore repr@(ArrayR shr _) exe gamma aenv halo shOut paramsR params =
withExecutable exe $ \ptxExecutable -> do
let
inside = ptxExecutable !# "stencil_inside"
border = ptxExecutable !# "stencil_border"
shIn :: sh
shIn = trav (\x y -> x - 2*y) shOut halo
trav :: (Int -> Int -> Int) -> sh -> sh -> sh
trav f a b = go (arrayRshape repr) a b
where
go :: ShapeR t -> t -> t -> t
go ShapeRz () () = ()
go (ShapeRsnoc shr') (xa,xb) (ya,yb) = (go shr' xa ya, f xb yb)
--
future <- new
result <- allocateRemote repr shOut
parent <- asks ptxStream
-- interior (no bounds checking)
let paramsRinside = TupRsingle (ParamRshape shr) `TupRpair` TupRsingle (ParamRarray repr) `TupRpair` paramsR
executeOp inside gamma aenv shr shIn paramsRinside ((shIn, result), params)
-- halo regions (bounds checking)
-- executed in separate streams so that they might overlap the main stencil
-- and each other, as individually they will not saturate the device
forM_ (stencilBorders (arrayRshape repr) shOut halo) $ \(u, v) ->
fork $ do
-- launch in a separate stream
let sh = trav (-) v u
let paramsRborder = TupRsingle (ParamRshape shr) `TupRpair` TupRsingle (ParamRshape shr)
`TupRpair` TupRsingle (ParamRarray repr)
`TupRpair` paramsR
executeOp border gamma aenv shr sh paramsRborder (((u, sh), result), params)
-- synchronisation with main stream
child <- asks ptxStream
event <- liftPar (Event.waypoint child)
ready <- liftIO (Event.query event)
if ready then return ()
else liftIO (Event.after event parent)
put future result
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
-> [(sh, sh)]
stencilBorders shr sh halo = [ face i | i <- [0 .. (2 * rank shr - 1)] ]
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'))
-- Foreign functions
--
{-# INLINE aforeignOp #-}
aforeignOp
:: HasCallStack
=> String
-> ArraysR as
-> ArraysR bs
-> (as -> Par PTX (Future bs))
-> as
-> Par PTX (Future bs)
aforeignOp name _ _ asm arr = do
stream <- asks ptxStream
Debug.monitorProcTime query msg (Just (unsafeGetValue stream)) (asm arr)
where
query = if Debug.monitoringIsEnabled
then return True
else liftIO $ Debug.getFlag Debug.dump_exec
msg wall cpu gpu = do
Debug.addProcessorTime Debug.PTX gpu
Debug.traceIO Debug.dump_exec $
printf "exec: %s %s" name (Debug.elapsed wall cpu gpu)
-- Skeleton execution
-- ------------------
-- | Retrieve the named kernel
--
(!#) :: HasCallStack => FunctionTable -> ShortByteString -> Kernel
(!#) exe name
= fromMaybe (internalError ("function not found: " ++ unpack name))
$ lookupKernel name exe
lookupKernel :: ShortByteString -> FunctionTable -> Maybe Kernel
lookupKernel name ptxExecutable =
find (\k -> kernelName k == name) (functionTable ptxExecutable)
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
-- | Execute some operation with the supplied executable functions
--
withExecutable :: HasCallStack => ExecutableR PTX -> (FunctionTable -> Par PTX b) -> Par PTX b
withExecutable PTXR{..} f =
local (\(s,_) -> (s,Just ptxExecutable)) $ do
r <- f (unsafeGetValue ptxExecutable)
liftIO $ touchLifetime ptxExecutable
return r
-- Execute the function implementing this kernel.
--
executeOp
:: HasCallStack
=> Kernel
-> Gamma aenv
-> Val aenv
-> ShapeR sh
-> sh
-> ParamsR PTX params
-> params
-> Par PTX ()
executeOp kernel gamma aenv shr sh paramsR params =
let n = size shr sh
in when (n > 0) $ do
stream <- asks ptxStream
argv <- marshalParams' @PTX (paramsR `TupRpair` TupRsingle (ParamRenv gamma)) (params, aenv)
liftIO $ launch kernel stream n $ DL.toList argv
-- Execute a device function with the given thread configuration and function
-- parameters.
--
launch :: HasCallStack => Kernel -> Stream -> Int -> [CUDA.FunParam] -> IO ()
launch Kernel{..} stream n args =
withLifetime stream $ \st ->
Debug.monitorProcTime query msg (Just st) $
CUDA.launchKernel kernelFun grid cta smem (Just st) args
where
cta = (kernelThreadBlockSize, 1, 1)
grid = (kernelThreadBlocks n, 1, 1)
smem = kernelSharedMemBytes
-- Debugging/monitoring support
query = if Debug.monitoringIsEnabled
then return True
else Debug.getFlag Debug.dump_exec
fst3 (x,_,_) = x
msg wall cpu gpu = do
Debug.addProcessorTime Debug.PTX gpu
Debug.traceIO Debug.dump_exec $
printf "exec: %s <<< %d, %d, %d >>> %s"
(unpack kernelName) (fst3 grid) (fst3 cta) smem (Debug.elapsed wall cpu gpu)