content: {-# 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
  
  