accelerate-llvm-native-1.4.0.0: src/Data/Array/Accelerate/LLVM/Native/CodeGen/Stencil.hs
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
-- |
-- Module : Data.Array.Accelerate.LLVM.Native.CodeGen.Stencil
-- Copyright : [2018..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.CodeGen.Stencil (
mkStencil1,
mkStencil2,
) where
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Shape
import Data.Array.Accelerate.Representation.Stencil
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Type
import Data.Array.Accelerate.LLVM.CodeGen.Arithmetic
import Data.Array.Accelerate.LLVM.CodeGen.Array
import Data.Array.Accelerate.LLVM.CodeGen.Base
import Data.Array.Accelerate.LLVM.CodeGen.Environment
import Data.Array.Accelerate.LLVM.CodeGen.Exp
import Data.Array.Accelerate.LLVM.CodeGen.IR
import Data.Array.Accelerate.LLVM.CodeGen.Loop hiding ( imapFromStepTo )
import Data.Array.Accelerate.LLVM.CodeGen.Monad
import Data.Array.Accelerate.LLVM.CodeGen.Stencil
import Data.Array.Accelerate.LLVM.CodeGen.Sugar
import Data.Array.Accelerate.LLVM.Compile.Cache
import Data.Array.Accelerate.LLVM.Native.CodeGen.Base
import Data.Array.Accelerate.LLVM.Native.CodeGen.Loop
import Data.Array.Accelerate.LLVM.Native.Target ( Native )
import qualified Data.Array.Accelerate.LLVM.Internal.LLVMPretty as LP
import Control.Monad
-- The stencil function is similar to a map, but has access to surrounding
-- elements as specified by the stencil pattern.
--
-- This generates two functions:
--
-- * stencil_inside: does not apply boundary conditions, assumes all element
-- accesses are valid
--
-- * stencil_border: applies boundary condition check to each array access
--
mkStencil1
:: UID
-> Gamma aenv
-> StencilR sh a stencil
-> TypeR b
-> IRFun1 Native aenv (stencil -> b)
-> IRBoundary Native aenv (Array sh a)
-> MIRDelayed Native aenv (Array sh a)
-> CodeGen Native (IROpenAcc Native aenv (Array sh b))
mkStencil1 uid aenv sr tp f bnd marr =
let (arrIn, paramIn) = delayedArray "in" marr
repr = ArrayR (stencilShapeR sr) tp
in (+++) <$> mkInside uid aenv repr (IRFun1 $ app1 f <=< stencilAccess sr Nothing arrIn) paramIn
<*> mkBorder uid aenv repr (IRFun1 $ app1 f <=< stencilAccess sr (Just bnd) arrIn) paramIn
mkStencil2
:: UID
-> Gamma aenv
-> StencilR sh a stencil1
-> StencilR sh b stencil2
-> TypeR c
-> IRFun2 Native aenv (stencil1 -> stencil2 -> c)
-> IRBoundary Native aenv (Array sh a)
-> MIRDelayed Native aenv (Array sh a)
-> IRBoundary Native aenv (Array sh b)
-> MIRDelayed Native aenv (Array sh b)
-> CodeGen Native (IROpenAcc Native aenv (Array sh c))
mkStencil2 uid aenv sr1 sr2 tp f bnd1 marr1 bnd2 marr2 =
let
(arrIn1, paramIn1) = delayedArray "in1" marr1
(arrIn2, paramIn2) = delayedArray "in2" marr2
repr = ArrayR (stencilShapeR sr1) tp
inside = IRFun1 $ \ix -> do
stencil1 <- stencilAccess sr1 Nothing arrIn1 ix
stencil2 <- stencilAccess sr2 Nothing arrIn2 ix
app2 f stencil1 stencil2
--
border = IRFun1 $ \ix -> do
stencil1 <- stencilAccess sr1 (Just bnd1) arrIn1 ix
stencil2 <- stencilAccess sr2 (Just bnd2) arrIn2 ix
app2 f stencil1 stencil2
in
(+++) <$> mkInside uid aenv repr inside (paramIn1 ++ paramIn2)
<*> mkBorder uid aenv repr border (paramIn1 ++ paramIn2)
mkInside
:: UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> IRFun1 Native aenv (sh -> e)
-> [LP.Typed LP.Ident]
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
mkInside uid aenv repr apply paramIn =
let
(start, end, paramGang) = gangParam (arrayRshape repr)
(arrOut, paramOut) = mutableArray repr "out"
paramEnv = envParam aenv
shOut = irArrayShape arrOut
in
makeOpenAcc uid "stencil_inside" (paramGang ++ paramOut ++ paramIn ++ paramEnv) $ do
imapNestFromToTile (arrayRshape repr) 4 start end shOut $ \ix i -> do
r <- app1 apply ix -- apply generator function
writeArray TypeInt arrOut i r -- store result
mkBorder
:: UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> IRFun1 Native aenv (sh -> e)
-> [LP.Typed LP.Ident]
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
mkBorder uid aenv repr apply paramIn =
let
(start, end, paramGang) = gangParam (arrayRshape repr)
(arrOut, paramOut) = mutableArray repr "out"
paramEnv = envParam aenv
shOut = irArrayShape arrOut
in
makeOpenAcc uid "stencil_border" (paramGang ++ paramOut ++ paramIn ++ paramEnv) $ do
imapNestFromTo (arrayRshape repr) start end shOut $ \ix i -> do
r <- app1 apply ix -- apply generator function
writeArray TypeInt arrOut i r -- store result
imapNestFromToTile
:: ShapeR sh
-> Int -- ^ unroll amount (tile height)
-> Operands sh -- ^ initial index (inclusive)
-> Operands sh -- ^ final index (exclusive)
-> Operands sh -- ^ total array extent
-> (Operands sh -> Operands Int -> CodeGen Native ()) -- ^ apply at each index
-> CodeGen Native ()
imapNestFromToTile shr unroll start end extent body =
go shr start end body'
where
body' ix = body ix =<< intOfIndex shr extent ix
go :: ShapeR t
-> Operands t
-> Operands t
-> (Operands t -> CodeGen Native ())
-> CodeGen Native ()
go ShapeRz OP_Unit OP_Unit k
= k OP_Unit
-- To correctly generate the unrolled loop nest we need to explicitly match
-- on the last two dimensions.
--
go (ShapeRsnoc (ShapeRsnoc ShapeRz)) (OP_Pair (OP_Pair OP_Unit sy) sx) (OP_Pair (OP_Pair OP_Unit ey) ex) k
= do
-- Tile the stencil operator in the xy-plane by unrolling in the
-- y-dimension and vectorising in the x-dimension.
--
sy' <- imapFromStepTo sy (liftInt unroll) ey $ \iy ->
imapFromTo sx ex $ \ix ->
forM_ [0 .. unroll-1] $ \n -> do
iy' <- add numType iy (liftInt n)
k (OP_Pair (OP_Pair OP_Unit iy') ix)
-- Take care of any remaining loop iterations in the y-dimension
--
_ <- imapFromTo sy' ey $ \iy ->
imapFromTo sx ex $ \ix ->
k (OP_Pair (OP_Pair OP_Unit iy) ix)
return ()
-- The 1- and 3+-dimensional cases can recurse normally
--
go (ShapeRsnoc shr') (OP_Pair ssh ssz) (OP_Pair esh esz) k
= go shr' ssh esh
$ \sz -> imapFromTo ssz esz
$ \i -> k (OP_Pair sz i)
imapFromStepTo
:: Operands Int
-> Operands Int
-> Operands Int
-> (Operands Int -> CodeGen Native ())
-> CodeGen Native (Operands Int)
imapFromStepTo start step end body =
let
incr i = add numType i step
test i = do i' <- incr i
lt singleType i' end
in
while (TupRsingle scalarTypeInt) test
(\i -> body i >> incr i)
start