content: {-# LANGUAGE GADTs               #-}
  {-# LANGUAGE OverloadedStrings   #-}
  {-# 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 Data.String
  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 <- freshLocalName
             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)
  
  