accelerate-llvm-1.3.0.0: src/Data/Array/Accelerate/LLVM/CodeGen/Stencil.hs
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_HADDOCK hide #-}
-- |
-- Module : Data.Array.Accelerate.LLVM.CodeGen.Stencil
-- Copyright : [2016..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.CodeGen.Stencil
where
import Data.Array.Accelerate.Error
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 ( ifThenElse )
import Data.Array.Accelerate.LLVM.CodeGen.Constant
import Data.Array.Accelerate.LLVM.CodeGen.IR
import Data.Array.Accelerate.LLVM.CodeGen.Monad
import Data.Array.Accelerate.LLVM.CodeGen.Sugar
import qualified Data.Array.Accelerate.LLVM.CodeGen.Arithmetic as A
import Control.Applicative
import Prelude
-- Stencil boundary conditions
--
data IRBoundary arch aenv t where
IRClamp :: IRBoundary arch aenv t
IRMirror :: IRBoundary arch aenv t
IRWrap :: IRBoundary arch aenv t
IRConstant :: Operands e -> IRBoundary arch aenv (Array sh e)
IRFunction :: IRFun1 arch aenv (sh -> e) -> IRBoundary arch aenv (Array sh e)
-- Generate the stencil pattern including boundary conditions
--
stencilAccess
:: HasCallStack
=> StencilR sh e stencil
-> Maybe (IRBoundary arch aenv (Array sh e))
-> IRDelayed arch aenv (Array sh e)
-> Operands sh
-> IRExp arch aenv stencil
stencilAccess sR mbndy arr =
case mbndy of
Nothing -> goR sR (inbounds arr)
Just bndy -> goR sR (bounded bndy arr)
where
-- Base cases, nothing interesting to do here since we know the lower
-- dimension is Z.
--
goR :: StencilR sh e stencil
-> (Operands sh -> IRExp arch aenv e)
-> Operands sh
-> IRExp arch aenv stencil
goR (StencilRunit3 _) rf ix
= let (z, i) = unindex ix
rf' d = do d' <- A.add numType i (int d)
rf (index z d')
in
tup3 <$> rf' (-1)
<*> rf ix
<*> rf' 1
goR (StencilRunit5 _) rf ix
= let (z, i) = unindex ix
rf' d = do d' <- A.add numType i (int d)
rf (index z d')
in
tup5 <$> rf' (-2)
<*> rf' (-1)
<*> rf ix
<*> rf' 1
<*> rf' 2
goR (StencilRunit7 _) rf ix
= let (z, i) = unindex ix
rf' d = do d' <- A.add numType i (int d)
rf (index z d')
in
tup7 <$> rf' (-3)
<*> rf' (-2)
<*> rf' (-1)
<*> rf ix
<*> rf' 1
<*> rf' 2
<*> rf' 3
goR (StencilRunit9 _) rf ix
= let (z, i) = unindex ix
rf' d = do d' <- A.add numType i (int d)
rf (index z d')
in
tup9 <$> rf' (-4)
<*> rf' (-3)
<*> rf' (-2)
<*> rf' (-1)
<*> rf ix
<*> rf' 1
<*> rf' 2
<*> rf' 3
<*> rf' 4
-- Recursive cases. Note that because the stencil pattern is defined with
-- a cons ordering, whereas shapes (indices) are defined as a snoc list,
-- when we recurse on the stencil structure we must manipulate the
-- _left-most_ index component
--
goR (StencilRtup3 s1 s2 s3) rf ix =
let shr = stencilShapeR s1
(i, ix') = uncons shr ix
rf' 0 ds = rf (cons shr i ds)
rf' d ds = do d' <- A.add numType i (int d)
rf (cons shr d' ds)
in
tup3 <$> goR s1 (rf' (-1)) ix'
<*> goR s2 (rf' 0) ix'
<*> goR s3 (rf' 1) ix'
goR (StencilRtup5 s1 s2 s3 s4 s5) rf ix =
let shr = stencilShapeR s1
(i, ix') = uncons shr ix
rf' 0 ds = rf (cons shr i ds)
rf' d ds = do d' <- A.add numType i (int d)
rf (cons shr d' ds)
in
tup5 <$> goR s1 (rf' (-2)) ix'
<*> goR s2 (rf' (-1)) ix'
<*> goR s3 (rf' 0) ix'
<*> goR s4 (rf' 1) ix'
<*> goR s5 (rf' 2) ix'
goR (StencilRtup7 s1 s2 s3 s4 s5 s6 s7) rf ix =
let shr = stencilShapeR s1
(i, ix') = uncons shr ix
rf' 0 ds = rf (cons shr i ds)
rf' d ds = do d' <- A.add numType i (int d)
rf (cons shr d' ds)
in
tup7 <$> goR s1 (rf' (-3)) ix'
<*> goR s2 (rf' (-2)) ix'
<*> goR s3 (rf' (-1)) ix'
<*> goR s4 (rf' 0) ix'
<*> goR s5 (rf' 1) ix'
<*> goR s6 (rf' 2) ix'
<*> goR s7 (rf' 3) ix'
goR (StencilRtup9 s1 s2 s3 s4 s5 s6 s7 s8 s9) rf ix =
let shr = stencilShapeR s1
(i, ix') = uncons shr ix
rf' 0 ds = rf (cons shr i ds)
rf' d ds = do d' <- A.add numType i (int d)
rf (cons shr d' ds)
in
tup9 <$> goR s1 (rf' (-4)) ix'
<*> goR s2 (rf' (-3)) ix'
<*> goR s3 (rf' (-2)) ix'
<*> goR s4 (rf' (-1)) ix'
<*> goR s5 (rf' 0) ix'
<*> goR s6 (rf' 1) ix'
<*> goR s7 (rf' 2) ix'
<*> goR s8 (rf' 3) ix'
<*> goR s9 (rf' 4) ix'
-- Do not apply any boundary conditions to the given index
--
inbounds
:: IRDelayed arch aenv (Array sh e)
-> Operands sh
-> IRExp arch aenv e
inbounds IRDelayed{..} ix =
app1 delayedIndex ix
-- Apply boundary conditions to the given index
--
bounded
:: forall sh e arch aenv. HasCallStack
=> IRBoundary arch aenv (Array sh e)
-> IRDelayed arch aenv (Array sh e)
-> Operands sh
-> IRExp arch aenv e
bounded bndy IRDelayed{..} ix = do
let
tp :: TypeR e -- GHC 8.4 needs this type annotation
ArrayR shr tp = delayedRepr
sh <- delayedExtent
case bndy of
IRConstant v ->
if ( tp, inside shr sh ix )
then app1 delayedIndex ix
else return v
IRFunction f ->
if ( tp, inside shr sh ix )
then app1 delayedIndex ix
else app1 f ix
_ -> do
ix' <- bound shr sh ix
v <- app1 delayedIndex ix'
return v
--
where
-- Return the index, updated to obey the given boundary conditions (clamp,
-- mirror, or wrap only).
--
bound :: ShapeR sh' -> Operands sh' -> Operands sh' -> CodeGen arch (Operands sh')
bound ShapeRz OP_Unit OP_Unit
= return OP_Unit
bound (ShapeRsnoc shr') (OP_Pair sh sz) (OP_Pair ih iz)
= do
ix' <- bound shr' sh ih
i' <- if ( TupRsingle scalarTypeInt
, A.lt (singleType :: SingleType Int) iz (int 0))
then
case bndy of
IRClamp -> return (int 0)
IRMirror -> A.negate numType iz
IRWrap -> A.add numType sz iz
_ -> internalError "unexpected boundary condition"
else
if ( TupRsingle scalarTypeInt
, A.gte (singleType :: SingleType Int) iz sz)
then
case bndy of
IRClamp -> A.sub numType sz (int 1)
IRWrap -> A.sub numType iz sz
IRMirror -> do
a <- A.sub numType iz sz
b <- A.add numType a (int 2)
c <- A.sub numType sz b
return c
_ -> internalError "unexpected boundary condition"
else
return iz
return $ OP_Pair ix' i'
-- Return whether the index is inside the bounds of the given shape
--
inside :: ShapeR sh' -> Operands sh' -> Operands sh' -> CodeGen arch (Operands Bool)
inside ShapeRz OP_Unit OP_Unit
= return (bool True)
inside (ShapeRsnoc shr') (OP_Pair sh sz) (OP_Pair ih iz)
= do
ifNext <- newBlock "inside.next"
ifExit <- newBlock "inside.exit"
_ <- beginBlock "inside.entry"
p <- A.lt (singleType :: SingleType Int) iz (int 0) `A.lor'`
A.gte (singleType :: SingleType Int) iz sz
eb <- cbr p ifExit ifNext
setBlock ifNext
nv <- inside shr' sh ih
nb <- br ifExit
setBlock ifExit
crit <- freshName
r <- phi1 ifExit crit [(boolean False, eb), (A.unbool nv, nb)]
return (OP_Bool r)
-- Utilities
-- ---------
int :: Int -> Operands Int
int x = constant (TupRsingle scalarTypeInt) x
bool :: Bool -> Operands Bool
bool = OP_Bool . boolean
unindex :: Operands (sh, Int) -> (Operands sh, Operands Int)
unindex (OP_Pair sh i) = (sh, i)
index :: Operands sh -> Operands Int -> Operands (sh, Int)
index sh i = OP_Pair sh i
tup3 :: Operands a -> Operands b -> Operands c -> Operands (Tup3 a b c)
tup3 a b c = OP_Pair (OP_Pair (OP_Pair OP_Unit a) b) c
tup5 :: Operands a -> Operands b -> Operands c -> Operands d -> Operands e -> Operands (Tup5 a b c d e)
tup5 a b c d e =
OP_Pair (OP_Pair (OP_Pair (OP_Pair (OP_Pair OP_Unit a) b) c) d) e
tup7 :: Operands a -> Operands b -> Operands c -> Operands d -> Operands e -> Operands f -> Operands g -> Operands (Tup7 a b c d e f g)
tup7 a b c d e f g =
OP_Pair (OP_Pair (OP_Pair (OP_Pair (OP_Pair (OP_Pair (OP_Pair OP_Unit a) b) c) d) e) f) g
tup9 :: Operands a -> Operands b -> Operands c -> Operands d -> Operands e -> Operands f -> Operands g -> Operands h -> Operands i -> Operands (Tup9 a b c d e f g h i)
tup9 a b c d e f g h i =
OP_Pair (OP_Pair (OP_Pair (OP_Pair (OP_Pair (OP_Pair (OP_Pair (OP_Pair (OP_Pair OP_Unit a) b) c) d) e) f) g) h) i
-- Add a _left-most_ dimension to a shape
--
cons :: ShapeR sh -> Operands Int -> Operands sh -> Operands (sh, Int)
cons ShapeRz ix OP_Unit = OP_Pair OP_Unit ix
cons (ShapeRsnoc shr) ix (OP_Pair sh sz) = OP_Pair (cons shr ix sh) sz
-- Remove the _left-most_ index to a shape, and return the remainder
--
uncons :: ShapeR sh -> Operands (sh, Int) -> (Operands Int, Operands sh)
uncons ShapeRz (OP_Pair OP_Unit v2) = (v2, OP_Unit)
uncons (ShapeRsnoc shr) (OP_Pair v1 v3)
= let (i, v1') = uncons shr v1
in (i, OP_Pair v1' v3)