accelerate-cuda-0.12.0.0: Data/Array/Accelerate/CUDA/CodeGen/Stencil.hs
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS -fno-warn-incomplete-patterns #-}
-- |
-- Module : Data.Array.Accelerate.CUDA.CodeGen.Mapping
-- Copyright : [2008..2010] Manuel M T Chakravarty, Gabriele Keller, Sean Lee
-- [2009..2012] Manuel M T Chakravarty, Gabriele Keller, Trevor L. McDonell
-- License : BSD3
--
-- Maintainer : Trevor L. McDonell <tmcdonell@cse.unsw.edu.au>
-- Stability : experimental
-- Portability : non-portable (GHC extensions)
--
module Data.Array.Accelerate.CUDA.CodeGen.Stencil (
mkStencil, mkStencil2
) where
import Language.C.Syntax
import Language.C.Quote.CUDA
import Data.Array.Accelerate.Type
import Data.Array.Accelerate.AST ( OpenAcc, Fun, Stencil )
import Data.Array.Accelerate.Array.Sugar ( Array, Elt, shapeToList )
import Data.Array.Accelerate.CUDA.CodeGen.Base
import Data.Array.Accelerate.CUDA.CodeGen.Type
import qualified Data.Array.Accelerate.Analysis.Stencil as Stencil
import qualified Data.Array.IArray as IArray
-- Map a stencil over an array. In contrast to 'map', the domain of a stencil
-- function is an entire /neighbourhood/ of each array element. Neighbourhoods
-- are sub-arrays centred around a focal point. They are not necessarily
-- rectangular, but they are symmetric in each dimension and have an extent of
-- at least three in each dimensions — due to the symmetry requirement, the
-- extent is necessarily odd. The focal point is the array position that is
-- determined by the stencil.
--
-- For those array positions where the neighbourhood extends past the boundaries
-- of the source array, a boundary condition determines the contents of the
-- out-of-bounds neighbourhood positions.
--
-- stencil :: (Shape ix, Elt a, Elt b, Stencil ix a stencil)
-- => (stencil -> Exp b) -- stencil function
-- -> Boundary a -- boundary condition
-- -> Acc (Array ix a) -- source array
-- -> Acc (Array ix b) -- destination array
--
-- To improve performance, the input array(s) are read through the texture
-- cache.
--
mkStencil :: forall sh stencil a b. (Stencil sh a stencil, Elt b)
=> Int
-> CUFun (stencil -> b)
-> Boundary (CUExp a)
-> Array sh b {- dummy -}
-> CUTranslSkel
mkStencil dim (CULam use0 (CUBody (CUExp env stencil))) boundary _ =
CUTranslSkel "stencil" [cunit|
$edecl:(cdim "Shape" dim)
$edecls:arrIn0
extern "C"
__global__ void
stencil
(
$params:argOut,
const typename Shape shIn0
)
{
const int shapeSize = size(shIn0);
const int gridSize = __umul24(blockDim.x, gridDim.x);
int i;
for ( i = __umul24(blockDim.x, blockIdx.x) + threadIdx.x
; i < shapeSize
; i += gridSize )
{
const typename Shape ix = fromIndex(shIn0, i);
$decls:(getIn0 "ix")
$decls:env
$stms:(setOut "i" stencil)
}
}
|]
where
tyOut = eltType (undefined :: b)
stencilIn0 = eltTypeTex (undefined :: a)
(argOut, _, setOut) = setters tyOut
(arrIn0, getIn0) = stencilAccess 0 dim stencilIn0 use0 boundary offsets
--
offsets = map shapeToList p0
p0 = Stencil.offsets (undefined :: Fun aenv (stencil -> b))
(undefined :: OpenAcc aenv (Array sh a))
-- Map a binary stencil of an array. The extent of the resulting array is the
-- intersection of the extents of the two source arrays.
--
-- stencil2 :: (Shape ix, Elt a, Elt b, Elt c,
-- Stencil ix a stencil1,
-- Stencil ix b stencil2)
-- => (stencil1 -> stencil2 -> Exp c) -- binary stencil function
-- -> Boundary a -- boundary condition #1
-- -> Acc (Array ix a) -- source array #1
-- -> Boundary b -- boundary condition #2
-- -> Acc (Array ix b) -- source array #2
-- -> Acc (Array ix c) -- destination array
--
mkStencil2 :: forall sh stencil1 stencil0 a b c.
(Stencil sh a stencil1, Stencil sh b stencil0, Elt c)
=> Int
-> CUFun (stencil1 -> stencil0 -> c)
-> Boundary (CUExp a)
-> Boundary (CUExp b)
-> Array sh c {- dummy -}
-> CUTranslSkel
mkStencil2 dim (CULam use1 (CULam use0 (CUBody (CUExp env stencil)))) boundary1 boundary0 _ =
CUTranslSkel "stencil2" [cunit|
$edecl:(cdim "Shape" dim)
$edecls:arrIn0
$edecls:arrIn1
extern "C"
__global__ void
stencil2
(
$params:argOut,
const typename Shape shOut,
const typename Shape shIn1,
const typename Shape shIn0
)
{
const int shapeSize = size(shOut);
const int gridSize = __umul24(blockDim.x, gridDim.x);
int i;
for ( i = __umul24(blockDim.x, blockIdx.x) + threadIdx.x
; i < shapeSize
; i += gridSize )
{
const typename Shape ix = fromIndex(shOut, i);
$decls:(getIn0 "ix")
$decls:(getIn1 "ix")
$decls:env
$stms:(setOut "i" stencil)
}
}
|]
where
tyOut = eltType (undefined :: c)
stencilIn0 = eltTypeTex (undefined :: b)
stencilIn1 = eltTypeTex (undefined :: a)
(argOut, _, setOut) = setters tyOut
(arrIn0, getIn0) = stencilAccess 0 dim stencilIn0 use0 boundary0 offsets0
(arrIn1, getIn1) = stencilAccess 1 dim stencilIn1 use1 boundary1 offsets1
--
offsets0 = map shapeToList p0
offsets1 = map shapeToList p1
(p1, p0) = Stencil.offsets2 (undefined :: Fun aenv (stencil1 -> stencil0 -> c))
(undefined :: OpenAcc aenv (Array sh a))
(undefined :: OpenAcc aenv (Array sh b))
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------
stencilAccess
:: Int -- array de Bruijn index
-> Int -- array dimensionality
-> [Type] -- array type (texture memory)
-> [(Int, Type, Exp)] -- the variables used in the scalar expression
-> Boundary (CUExp a) -- how to handle boundary array access
-> [[Int]] -- all stencil index offsets, top left to bottom right
-> ( [Definition] -- texture-reference definitions
, String -> [InitGroup] ) -- array indexing
stencilAccess base dim stencil subs boundary shx =
( textures
, \ix -> concatMap (get ix) subs )
where
n = length stencil
sh = "shIn" ++ show base
arr x = "stencil" ++ shows base "_a" ++ show (x `mod` n)
textures = zipWith cglobal stencil (map arr [n-1, n-2 .. 0])
--
offsets :: IArray.Array Int [Int]
offsets = IArray.listArray (0, length shx-1) shx
--
get ix (i,t,v) = case boundary of
Clamp -> bounded "clamp"
Mirror -> bounded "mirror"
Wrap -> bounded "wrap"
Constant (CUExp _ c) -> inRange c
where
j = 'j':shows base "_a" ++ show i
k = 'k':shows base "_a" ++ show i
--
bounded f
= [cdecl| const int $id:j = $exp:ix'; |]
: [cdecl| const $ty:t $id:(show v) = $exp:(indexArray t (cvar (arr i)) (cvar j)); |]
: []
where
ix' = case offsets IArray.! div i n of
ks | all (== 0) ks -> [cexp| toIndex( $id:sh, ix ) |]
| otherwise -> [cexp| toIndex( $id:sh, $exp:(ccall f [cvar sh, cursor ks]) ) |]
--
inRange c = case offsets IArray.! div i n of
ks | all (== 0) ks -> let f = indexArray t (cvar (arr i)) (ccall "toIndex" [cvar sh, cvar "ix"])
in [[cdecl| const $ty:t $id:(show v) = $exp:f; |]]
| otherwise -> [cdecl| const typename Shape $id:j = $exp:(cursor ks); |]
: [cdecl| const typename bool $id:k = inRange( $id:sh, $id:j ); |]
: [cdecl| const $ty:t $id:(show v) = $id:k ? $exp:(indexArray t (cvar (arr i)) (ccall "toIndex" [cvar sh, cvar j]))
: $exp:(reverse c !! mod i n); |]
: []
--
cursor [c] = [cexp| $id:ix + $int:c |]
cursor cs = ccall "shape" $ zipWith (\a c -> [cexp| $id:ix . $id:('a':show a) + $int:c |]) [dim-1,dim-2..0] cs