Stencil.hs

{-# LANGUAGE GADTs               #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeApplications    #-}
-- |
-- Module      : Data.Array.Accelerate.LLVM.PTX.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.PTX.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.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.PTX.CodeGen.Base
import Data.Array.Accelerate.LLVM.PTX.CodeGen.Loop
import Data.Array.Accelerate.LLVM.PTX.Target                        ( PTX )

import qualified Data.Array.Accelerate.LLVM.Internal.LLVMPretty.AST 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      PTX aenv (stencil -> b)
    -> IRBoundary  PTX aenv (Array sh a)
    -> MIRDelayed  PTX aenv (Array sh a)
    -> CodeGen     PTX      (IROpenAcc PTX aenv (Array sh b))
mkStencil1 uid aenv stencil tp fun bnd marr =
  let repr             = ArrayR shr tp
      (shr, halo)      = stencilHalo stencil
      (arrIn, paramIn) = delayedArray "in" marr
  in
  (+++) <$> mkInside uid aenv repr halo (IRFun1 $ app1 fun <=< stencilAccess stencil Nothing    arrIn) paramIn
        <*> mkBorder uid aenv repr      (IRFun1 $ app1 fun <=< stencilAccess stencil (Just bnd) arrIn) paramIn


mkStencil2
    :: UID
    -> Gamma           aenv
    -> StencilR sh a stencil1
    -> StencilR sh b stencil2
    -> TypeR c
    -> IRFun2      PTX aenv (stencil1 -> stencil2 -> c)
    -> IRBoundary  PTX aenv (Array sh a)
    -> MIRDelayed  PTX aenv (Array sh a)
    -> IRBoundary  PTX aenv (Array sh b)
    -> MIRDelayed  PTX aenv (Array sh b)
    -> CodeGen     PTX      (IROpenAcc PTX aenv (Array sh c))
mkStencil2 uid aenv stencil1 stencil2 tp f bnd1 marr1 bnd2 marr2 =
  let
      repr = ArrayR shr tp
      (arrIn1, paramIn1)  = delayedArray "in1" marr1
      (arrIn2, paramIn2)  = delayedArray "in2" marr2

      inside  = IRFun1 $ \ix -> do
        s1 <- stencilAccess stencil1 Nothing arrIn1 ix
        s2 <- stencilAccess stencil2 Nothing arrIn2 ix
        app2 f s1 s2
      --
      border  = IRFun1 $ \ix -> do
        s1 <- stencilAccess stencil1 (Just bnd1) arrIn1 ix
        s2 <- stencilAccess stencil2 (Just bnd2) arrIn2 ix
        app2 f s1 s2

      (shr, halo1) = stencilHalo stencil1
      (_,   halo2) = stencilHalo stencil2
      halo         = union shr halo1 halo2
  in
  (+++) <$> mkInside uid aenv repr halo inside (paramIn1 ++ paramIn2)
        <*> mkBorder uid aenv repr      border (paramIn1 ++ paramIn2)


mkInside
    :: UID
    -> Gamma aenv
    -> ArrayR (Array sh e)
    -> sh
    -> IRFun1  PTX aenv (sh -> e)
    -> [LP.Typed LP.Ident]
    -> CodeGen PTX      (IROpenAcc PTX aenv (Array sh e))
mkInside uid aenv repr@(ArrayR shr _) halo apply paramIn =
  let
      (arrOut, paramOut)  = mutableArray repr "out"
      paramInside         = parameter    (shapeType shr) "shInside"
      shInside            = local        (shapeType shr) "shInside"
      shOut               = irArrayShape arrOut
      paramEnv            = envParam aenv
  in
  makeOpenAcc uid "stencil_inside" (paramInside ++ paramOut ++ paramIn ++ paramEnv) $ do

    start <- return (liftInt 0)
    end   <- shapeSize shr shInside

    -- iterate over the inside region as a linear index space
    --
    imapFromTo start end $ \i -> do

      ixIn  <- indexOfInt shr shInside i                    -- convert to multidimensional index of inside region
      ixOut <- offset shr ixIn (lift (shapeType shr) halo)  -- shift to multidimensional index of outside region
      r     <- app1 apply ixOut                             -- apply generator function
      j     <- intOfIndex shr shOut ixOut
      writeArray TypeInt arrOut j r

    return_


mkBorder
    :: UID
    -> Gamma aenv
    -> ArrayR (Array sh e)
    -> IRFun1  PTX aenv (sh -> e)
    -> [LP.Typed LP.Ident]
    -> CodeGen PTX      (IROpenAcc PTX aenv (Array sh e))
mkBorder uid aenv repr@(ArrayR shr _) apply paramIn =
  let
      (arrOut, paramOut)  = mutableArray repr "out"
      paramFrom           = parameter    (shapeType shr) "shFrom"
      shFrom              = local        (shapeType shr) "shFrom"
      paramInside         = parameter    (shapeType shr) "shInside"
      shInside            = local        (shapeType shr) "shInside"
      shOut               = irArrayShape arrOut
      paramEnv            = envParam aenv
  in
  makeOpenAcc uid "stencil_border" (paramFrom ++ paramInside ++ paramOut ++ paramIn ++ paramEnv) $ do

    start <- return (liftInt 0)
    end   <- shapeSize shr shInside

    imapFromTo start end $ \i -> do

      ixIn  <- indexOfInt shr shInside i    -- convert to multidimensional index of inside region
      ixOut <- offset shr ixIn shFrom       -- shift to multidimensional index of outside region
      r     <- app1 apply ixOut             -- apply generator function
      j     <- intOfIndex shr shOut ixOut
      writeArray TypeInt arrOut j r

    return_


offset :: ShapeR sh -> Operands sh -> Operands sh -> CodeGen PTX (Operands sh)
offset shr sh1 sh2 = go shr sh1 sh2
  where
    go :: ShapeR t -> Operands t -> Operands t -> CodeGen PTX (Operands t)
    go ShapeRz OP_Unit OP_Unit
      = return OP_Unit

    go (ShapeRsnoc t) (OP_Pair sa1 sb1) (OP_Pair sa2 sb2)
      = do x <- add (numType :: NumType Int) sb1 sb2
           OP_Pair <$> go t sa1 sa2 <*> return x