{-# 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.PTX.CodeGen.Base
import Data.Array.Accelerate.LLVM.PTX.CodeGen.Loop
import Data.Array.Accelerate.LLVM.PTX.Target                        ( PTX )

import qualified LLVM.AST.Global                                    as LLVM

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
    :: 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 :: 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 Gamma aenv
aenv StencilR sh a stencil
stencil TypeR b
tp IRFun1 PTX aenv (stencil -> b)
fun IRBoundary PTX aenv (Array sh a)
bnd MIRDelayed PTX aenv (Array sh a)
marr =
  let repr :: ArrayR (Array sh b)
repr             = ShapeR sh -> TypeR b -> ArrayR (Array sh b)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
shr TypeR b
tp
      (ShapeR sh
shr, sh
halo)      = StencilR sh a stencil -> (ShapeR sh, sh)
forall sh e stencil. StencilR sh e stencil -> (ShapeR sh, sh)
stencilHalo StencilR sh a stencil
stencil
      (IRDelayed PTX aenv (Array sh a)
arrIn, [Parameter]
paramIn) = Name (Array sh a)
-> MIRDelayed PTX aenv (Array sh a)
-> (IRDelayed PTX aenv (Array sh a), [Parameter])
forall sh e arch aenv.
Name (Array sh e)
-> MIRDelayed arch aenv (Array sh e)
-> (IRDelayed arch aenv (Array sh e), [Parameter])
delayedArray Name (Array sh a)
"in" MIRDelayed PTX aenv (Array sh a)
marr
  in
  IROpenAcc PTX aenv (Array sh b)
-> IROpenAcc PTX aenv (Array sh b)
-> IROpenAcc PTX aenv (Array sh b)
forall aenv a.
IROpenAcc PTX aenv a
-> IROpenAcc PTX aenv a -> IROpenAcc PTX aenv a
(+++) (IROpenAcc PTX aenv (Array sh b)
 -> IROpenAcc PTX aenv (Array sh b)
 -> IROpenAcc PTX aenv (Array sh b))
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh b))
-> CodeGen
     PTX
     (IROpenAcc PTX aenv (Array sh b)
      -> IROpenAcc PTX aenv (Array sh b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gamma aenv
-> ArrayR (Array sh b)
-> sh
-> IRFun1 PTX aenv (sh -> b)
-> [Parameter]
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh b))
forall aenv sh e.
Gamma aenv
-> ArrayR (Array sh e)
-> sh
-> IRFun1 PTX aenv (sh -> e)
-> [Parameter]
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh e))
mkInside Gamma aenv
aenv ArrayR (Array sh b)
repr sh
halo ((Operands sh -> IROpenExp PTX ((), sh) aenv b)
-> IRFun1 PTX aenv (sh -> b)
forall a arch env aenv b.
(Operands a -> IROpenExp arch (env, a) aenv b)
-> IROpenFun1 arch env aenv (a -> b)
IRFun1 ((Operands sh -> IROpenExp PTX ((), sh) aenv b)
 -> IRFun1 PTX aenv (sh -> b))
-> (Operands sh -> IROpenExp PTX ((), sh) aenv b)
-> IRFun1 PTX aenv (sh -> b)
forall a b. (a -> b) -> a -> b
$ IRFun1 PTX aenv (stencil -> b)
-> Operands stencil -> IROpenExp PTX ((), sh) aenv b
forall arch env aenv a b.
IROpenFun1 arch env aenv (a -> b)
-> Operands a -> IROpenExp arch (env, a) aenv b
app1 IRFun1 PTX aenv (stencil -> b)
fun (Operands stencil -> IROpenExp PTX ((), sh) aenv b)
-> (Operands sh -> CodeGen PTX (Operands stencil))
-> Operands sh
-> IROpenExp PTX ((), sh) aenv b
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< StencilR sh a stencil
-> Maybe (IRBoundary PTX aenv (Array sh a))
-> IRDelayed PTX aenv (Array sh a)
-> Operands sh
-> CodeGen PTX (Operands stencil)
forall sh e stencil arch aenv.
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 StencilR sh a stencil
stencil Maybe (IRBoundary PTX aenv (Array sh a))
forall a. Maybe a
Nothing    IRDelayed PTX aenv (Array sh a)
arrIn) [Parameter]
paramIn
        CodeGen
  PTX
  (IROpenAcc PTX aenv (Array sh b)
   -> IROpenAcc PTX aenv (Array sh b))
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh b))
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh b))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gamma aenv
-> ArrayR (Array sh b)
-> IRFun1 PTX aenv (sh -> b)
-> [Parameter]
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh b))
forall aenv sh e.
Gamma aenv
-> ArrayR (Array sh e)
-> IRFun1 PTX aenv (sh -> e)
-> [Parameter]
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh e))
mkBorder Gamma aenv
aenv ArrayR (Array sh b)
repr      ((Operands sh -> IROpenExp PTX ((), sh) aenv b)
-> IRFun1 PTX aenv (sh -> b)
forall a arch env aenv b.
(Operands a -> IROpenExp arch (env, a) aenv b)
-> IROpenFun1 arch env aenv (a -> b)
IRFun1 ((Operands sh -> IROpenExp PTX ((), sh) aenv b)
 -> IRFun1 PTX aenv (sh -> b))
-> (Operands sh -> IROpenExp PTX ((), sh) aenv b)
-> IRFun1 PTX aenv (sh -> b)
forall a b. (a -> b) -> a -> b
$ IRFun1 PTX aenv (stencil -> b)
-> Operands stencil -> IROpenExp PTX ((), sh) aenv b
forall arch env aenv a b.
IROpenFun1 arch env aenv (a -> b)
-> Operands a -> IROpenExp arch (env, a) aenv b
app1 IRFun1 PTX aenv (stencil -> b)
fun (Operands stencil -> IROpenExp PTX ((), sh) aenv b)
-> (Operands sh -> CodeGen PTX (Operands stencil))
-> Operands sh
-> IROpenExp PTX ((), sh) aenv b
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< StencilR sh a stencil
-> Maybe (IRBoundary PTX aenv (Array sh a))
-> IRDelayed PTX aenv (Array sh a)
-> Operands sh
-> CodeGen PTX (Operands stencil)
forall sh e stencil arch aenv.
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 StencilR sh a stencil
stencil (IRBoundary PTX aenv (Array sh a)
-> Maybe (IRBoundary PTX aenv (Array sh a))
forall a. a -> Maybe a
Just IRBoundary PTX aenv (Array sh a)
bnd) IRDelayed PTX aenv (Array sh a)
arrIn) [Parameter]
paramIn


mkStencil2
    :: 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 :: 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 Gamma aenv
aenv StencilR sh a stencil1
stencil1 StencilR sh b stencil2
stencil2 TypeR c
tp IRFun2 PTX aenv (stencil1 -> stencil2 -> c)
f IRBoundary PTX aenv (Array sh a)
bnd1 MIRDelayed PTX aenv (Array sh a)
marr1 IRBoundary PTX aenv (Array sh b)
bnd2 MIRDelayed PTX aenv (Array sh b)
marr2 =
  let
      repr :: ArrayR (Array sh c)
repr = ShapeR sh -> TypeR c -> ArrayR (Array sh c)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
shr TypeR c
tp
      (IRDelayed PTX aenv (Array sh a)
arrIn1, [Parameter]
paramIn1)  = Name (Array sh a)
-> MIRDelayed PTX aenv (Array sh a)
-> (IRDelayed PTX aenv (Array sh a), [Parameter])
forall sh e arch aenv.
Name (Array sh e)
-> MIRDelayed arch aenv (Array sh e)
-> (IRDelayed arch aenv (Array sh e), [Parameter])
delayedArray Name (Array sh a)
"in1" MIRDelayed PTX aenv (Array sh a)
marr1
      (IRDelayed PTX aenv (Array sh b)
arrIn2, [Parameter]
paramIn2)  = Name (Array sh b)
-> MIRDelayed PTX aenv (Array sh b)
-> (IRDelayed PTX aenv (Array sh b), [Parameter])
forall sh e arch aenv.
Name (Array sh e)
-> MIRDelayed arch aenv (Array sh e)
-> (IRDelayed arch aenv (Array sh e), [Parameter])
delayedArray Name (Array sh b)
"in2" MIRDelayed PTX aenv (Array sh b)
marr2

      inside :: IROpenFun1 PTX () aenv (sh -> c)
inside  = (Operands sh -> IROpenExp PTX ((), sh) aenv c)
-> IROpenFun1 PTX () aenv (sh -> c)
forall a arch env aenv b.
(Operands a -> IROpenExp arch (env, a) aenv b)
-> IROpenFun1 arch env aenv (a -> b)
IRFun1 ((Operands sh -> IROpenExp PTX ((), sh) aenv c)
 -> IROpenFun1 PTX () aenv (sh -> c))
-> (Operands sh -> IROpenExp PTX ((), sh) aenv c)
-> IROpenFun1 PTX () aenv (sh -> c)
forall a b. (a -> b) -> a -> b
$ \Operands sh
ix -> do
        Operands stencil1
s1 <- StencilR sh a stencil1
-> Maybe (IRBoundary PTX aenv (Array sh a))
-> IRDelayed PTX aenv (Array sh a)
-> Operands sh
-> IRExp PTX aenv stencil1
forall sh e stencil arch aenv.
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 StencilR sh a stencil1
stencil1 Maybe (IRBoundary PTX aenv (Array sh a))
forall a. Maybe a
Nothing IRDelayed PTX aenv (Array sh a)
arrIn1 Operands sh
ix
        Operands stencil2
s2 <- StencilR sh b stencil2
-> Maybe (IRBoundary PTX aenv (Array sh b))
-> IRDelayed PTX aenv (Array sh b)
-> Operands sh
-> IRExp PTX aenv stencil2
forall sh e stencil arch aenv.
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 StencilR sh b stencil2
stencil2 Maybe (IRBoundary PTX aenv (Array sh b))
forall a. Maybe a
Nothing IRDelayed PTX aenv (Array sh b)
arrIn2 Operands sh
ix
        IRFun2 PTX aenv (stencil1 -> stencil2 -> c)
-> Operands stencil1
-> Operands stencil2
-> IROpenExp PTX ((), sh) aenv c
forall arch env aenv a b c.
IROpenFun2 arch env aenv (a -> b -> c)
-> Operands a -> Operands b -> IROpenExp arch ((env, a), b) aenv c
app2 IRFun2 PTX aenv (stencil1 -> stencil2 -> c)
f Operands stencil1
s1 Operands stencil2
s2
      --
      border :: IROpenFun1 PTX () aenv (sh -> c)
border  = (Operands sh -> IROpenExp PTX ((), sh) aenv c)
-> IROpenFun1 PTX () aenv (sh -> c)
forall a arch env aenv b.
(Operands a -> IROpenExp arch (env, a) aenv b)
-> IROpenFun1 arch env aenv (a -> b)
IRFun1 ((Operands sh -> IROpenExp PTX ((), sh) aenv c)
 -> IROpenFun1 PTX () aenv (sh -> c))
-> (Operands sh -> IROpenExp PTX ((), sh) aenv c)
-> IROpenFun1 PTX () aenv (sh -> c)
forall a b. (a -> b) -> a -> b
$ \Operands sh
ix -> do
        Operands stencil1
s1 <- StencilR sh a stencil1
-> Maybe (IRBoundary PTX aenv (Array sh a))
-> IRDelayed PTX aenv (Array sh a)
-> Operands sh
-> IRExp PTX aenv stencil1
forall sh e stencil arch aenv.
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 StencilR sh a stencil1
stencil1 (IRBoundary PTX aenv (Array sh a)
-> Maybe (IRBoundary PTX aenv (Array sh a))
forall a. a -> Maybe a
Just IRBoundary PTX aenv (Array sh a)
bnd1) IRDelayed PTX aenv (Array sh a)
arrIn1 Operands sh
ix
        Operands stencil2
s2 <- StencilR sh b stencil2
-> Maybe (IRBoundary PTX aenv (Array sh b))
-> IRDelayed PTX aenv (Array sh b)
-> Operands sh
-> IRExp PTX aenv stencil2
forall sh e stencil arch aenv.
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 StencilR sh b stencil2
stencil2 (IRBoundary PTX aenv (Array sh b)
-> Maybe (IRBoundary PTX aenv (Array sh b))
forall a. a -> Maybe a
Just IRBoundary PTX aenv (Array sh b)
bnd2) IRDelayed PTX aenv (Array sh b)
arrIn2 Operands sh
ix
        IRFun2 PTX aenv (stencil1 -> stencil2 -> c)
-> Operands stencil1
-> Operands stencil2
-> IROpenExp PTX ((), sh) aenv c
forall arch env aenv a b c.
IROpenFun2 arch env aenv (a -> b -> c)
-> Operands a -> Operands b -> IROpenExp arch ((env, a), b) aenv c
app2 IRFun2 PTX aenv (stencil1 -> stencil2 -> c)
f Operands stencil1
s1 Operands stencil2
s2

      (ShapeR sh
shr, sh
halo1) = StencilR sh a stencil1 -> (ShapeR sh, sh)
forall sh e stencil. StencilR sh e stencil -> (ShapeR sh, sh)
stencilHalo StencilR sh a stencil1
stencil1
      (ShapeR sh
_,   sh
halo2) = StencilR sh b stencil2 -> (ShapeR sh, sh)
forall sh e stencil. StencilR sh e stencil -> (ShapeR sh, sh)
stencilHalo StencilR sh b stencil2
stencil2
      halo :: sh
halo         = ShapeR sh -> sh -> sh -> sh
forall sh. ShapeR sh -> sh -> sh -> sh
union ShapeR sh
shr sh
halo1 sh
halo2
  in
  IROpenAcc PTX aenv (Array sh c)
-> IROpenAcc PTX aenv (Array sh c)
-> IROpenAcc PTX aenv (Array sh c)
forall aenv a.
IROpenAcc PTX aenv a
-> IROpenAcc PTX aenv a -> IROpenAcc PTX aenv a
(+++) (IROpenAcc PTX aenv (Array sh c)
 -> IROpenAcc PTX aenv (Array sh c)
 -> IROpenAcc PTX aenv (Array sh c))
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh c))
-> CodeGen
     PTX
     (IROpenAcc PTX aenv (Array sh c)
      -> IROpenAcc PTX aenv (Array sh c))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gamma aenv
-> ArrayR (Array sh c)
-> sh
-> IROpenFun1 PTX () aenv (sh -> c)
-> [Parameter]
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh c))
forall aenv sh e.
Gamma aenv
-> ArrayR (Array sh e)
-> sh
-> IRFun1 PTX aenv (sh -> e)
-> [Parameter]
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh e))
mkInside Gamma aenv
aenv ArrayR (Array sh c)
repr sh
halo IROpenFun1 PTX () aenv (sh -> c)
inside ([Parameter]
paramIn1 [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramIn2)
        CodeGen
  PTX
  (IROpenAcc PTX aenv (Array sh c)
   -> IROpenAcc PTX aenv (Array sh c))
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh c))
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh c))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gamma aenv
-> ArrayR (Array sh c)
-> IROpenFun1 PTX () aenv (sh -> c)
-> [Parameter]
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh c))
forall aenv sh e.
Gamma aenv
-> ArrayR (Array sh e)
-> IRFun1 PTX aenv (sh -> e)
-> [Parameter]
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh e))
mkBorder Gamma aenv
aenv ArrayR (Array sh c)
repr      IROpenFun1 PTX () aenv (sh -> c)
border ([Parameter]
paramIn1 [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramIn2)


mkInside
    :: Gamma aenv
    -> ArrayR (Array sh e)
    -> sh
    -> IRFun1  PTX aenv (sh -> e)
    -> [LLVM.Parameter]
    -> CodeGen PTX      (IROpenAcc PTX aenv (Array sh e))
mkInside :: Gamma aenv
-> ArrayR (Array sh e)
-> sh
-> IRFun1 PTX aenv (sh -> e)
-> [Parameter]
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh e))
mkInside Gamma aenv
aenv repr :: ArrayR (Array sh e)
repr@(ArrayR ShapeR sh
shr TypeR e
_) sh
halo IRFun1 PTX aenv (sh -> e)
apply [Parameter]
paramIn =
  let
      (IRArray (Array sh e)
arrOut, [Parameter]
paramOut)  = ArrayR (Array sh e)
-> Name (Array sh e) -> (IRArray (Array sh e), [Parameter])
forall sh e.
ArrayR (Array sh e)
-> Name (Array sh e) -> (IRArray (Array sh e), [Parameter])
mutableArray ArrayR (Array sh e)
repr Name (Array sh e)
"out"
      paramInside :: [Parameter]
paramInside         = TypeR sh -> Name sh -> [Parameter]
forall t. TypeR t -> Name t -> [Parameter]
parameter    (ShapeR sh -> TypeR sh
forall sh. ShapeR sh -> TypeR sh
shapeType ShapeR sh
shr) Name sh
"shInside"
      shInside :: Operands sh
shInside            = TypeR sh -> Name sh -> Operands sh
forall a. TypeR a -> Name a -> Operands a
local        (ShapeR sh -> TypeR sh
forall sh. ShapeR sh -> TypeR sh
shapeType ShapeR sh
shr) Name sh
"shInside"
      shOut :: Operands sh
shOut               = IRArray (Array sh e) -> Operands sh
forall sh e. IRArray (Array sh e) -> Operands sh
irArrayShape IRArray (Array sh e)
arrOut
      paramEnv :: [Parameter]
paramEnv            = Gamma aenv -> [Parameter]
forall aenv. Gamma aenv -> [Parameter]
envParam Gamma aenv
aenv
      --
  in
  Label
-> [Parameter]
-> CodeGen PTX ()
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh e))
forall aenv a.
Label
-> [Parameter]
-> CodeGen PTX ()
-> CodeGen PTX (IROpenAcc PTX aenv a)
makeOpenAcc Label
"stencil_inside" ([Parameter]
paramInside [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramOut [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramIn [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramEnv) (CodeGen PTX () -> CodeGen PTX (IROpenAcc PTX aenv (Array sh e)))
-> CodeGen PTX () -> CodeGen PTX (IROpenAcc PTX aenv (Array sh e))
forall a b. (a -> b) -> a -> b
$ do

    Operands Int
start <- Operands Int -> CodeGen PTX (Operands Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Operands Int
liftInt Int
0)
    Operands Int
end   <- ShapeR sh -> Operands sh -> CodeGen PTX (Operands Int)
forall sh arch.
ShapeR sh -> Operands sh -> CodeGen arch (Operands Int)
shapeSize ShapeR sh
shr Operands sh
shInside

    -- iterate over the inside region as a linear index space
    --
    Operands Int
-> Operands Int
-> (Operands Int -> CodeGen PTX ())
-> CodeGen PTX ()
imapFromTo Operands Int
start Operands Int
end ((Operands Int -> CodeGen PTX ()) -> CodeGen PTX ())
-> (Operands Int -> CodeGen PTX ()) -> CodeGen PTX ()
forall a b. (a -> b) -> a -> b
$ \Operands Int
i -> do

      Operands sh
ixIn  <- ShapeR sh
-> Operands sh -> Operands Int -> CodeGen PTX (Operands sh)
forall sh arch.
ShapeR sh
-> Operands sh -> Operands Int -> CodeGen arch (Operands sh)
indexOfInt ShapeR sh
shr Operands sh
shInside Operands Int
i                    -- convert to multidimensional index of inside region
      Operands sh
ixOut <- ShapeR sh
-> Operands sh -> Operands sh -> CodeGen PTX (Operands sh)
forall sh.
ShapeR sh
-> Operands sh -> Operands sh -> CodeGen PTX (Operands sh)
offset ShapeR sh
shr Operands sh
ixIn (TypeR sh -> sh -> Operands sh
forall a. TypeR a -> a -> Operands a
lift (ShapeR sh -> TypeR sh
forall sh. ShapeR sh -> TypeR sh
shapeType ShapeR sh
shr) sh
sh
halo)  -- shift to multidimensional index of outside region
      Operands e
r     <- IRFun1 PTX aenv (sh -> e)
-> Operands sh -> IROpenExp PTX ((), sh) aenv e
forall arch env aenv a b.
IROpenFun1 arch env aenv (a -> b)
-> Operands a -> IROpenExp arch (env, a) aenv b
app1 IRFun1 PTX aenv (sh -> e)
apply Operands sh
Operands sh
ixOut                             -- apply generator function
      Operands Int
j     <- ShapeR sh
-> Operands sh -> Operands sh -> CodeGen PTX (Operands Int)
forall sh arch.
ShapeR sh
-> Operands sh -> Operands sh -> CodeGen arch (Operands Int)
intOfIndex ShapeR sh
shr Operands sh
Operands sh
shOut Operands sh
ixOut
      IntegralType Int
-> IRArray (Array sh e)
-> Operands Int
-> Operands e
-> CodeGen PTX ()
forall int sh e arch.
IntegralType int
-> IRArray (Array sh e)
-> Operands int
-> Operands e
-> CodeGen arch ()
writeArray IntegralType Int
TypeInt IRArray (Array sh e)
arrOut Operands Int
j Operands e
r

    CodeGen PTX ()
forall arch. HasCallStack => CodeGen arch ()
return_


mkBorder
    :: Gamma aenv
    -> ArrayR (Array sh e)
    -> IRFun1  PTX aenv (sh -> e)
    -> [LLVM.Parameter]
    -> CodeGen PTX      (IROpenAcc PTX aenv (Array sh e))
mkBorder :: Gamma aenv
-> ArrayR (Array sh e)
-> IRFun1 PTX aenv (sh -> e)
-> [Parameter]
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh e))
mkBorder Gamma aenv
aenv repr :: ArrayR (Array sh e)
repr@(ArrayR ShapeR sh
shr TypeR e
_) IRFun1 PTX aenv (sh -> e)
apply [Parameter]
paramIn =
  let
      (IRArray (Array sh e)
arrOut, [Parameter]
paramOut)  = ArrayR (Array sh e)
-> Name (Array sh e) -> (IRArray (Array sh e), [Parameter])
forall sh e.
ArrayR (Array sh e)
-> Name (Array sh e) -> (IRArray (Array sh e), [Parameter])
mutableArray ArrayR (Array sh e)
repr Name (Array sh e)
"out"
      paramFrom :: [Parameter]
paramFrom           = TypeR sh -> Name sh -> [Parameter]
forall t. TypeR t -> Name t -> [Parameter]
parameter    (ShapeR sh -> TypeR sh
forall sh. ShapeR sh -> TypeR sh
shapeType ShapeR sh
shr) Name sh
"shFrom"
      shFrom :: Operands sh
shFrom              = TypeR sh -> Name sh -> Operands sh
forall a. TypeR a -> Name a -> Operands a
local        (ShapeR sh -> TypeR sh
forall sh. ShapeR sh -> TypeR sh
shapeType ShapeR sh
shr) Name sh
"shFrom"
      paramInside :: [Parameter]
paramInside         = TypeR sh -> Name sh -> [Parameter]
forall t. TypeR t -> Name t -> [Parameter]
parameter    (ShapeR sh -> TypeR sh
forall sh. ShapeR sh -> TypeR sh
shapeType ShapeR sh
shr) Name sh
"shInside"
      shInside :: Operands sh
shInside            = TypeR sh -> Name sh -> Operands sh
forall a. TypeR a -> Name a -> Operands a
local        (ShapeR sh -> TypeR sh
forall sh. ShapeR sh -> TypeR sh
shapeType ShapeR sh
shr) Name sh
"shInside"
      shOut :: Operands sh
shOut               = IRArray (Array sh e) -> Operands sh
forall sh e. IRArray (Array sh e) -> Operands sh
irArrayShape IRArray (Array sh e)
arrOut
      paramEnv :: [Parameter]
paramEnv            = Gamma aenv -> [Parameter]
forall aenv. Gamma aenv -> [Parameter]
envParam Gamma aenv
aenv
      --
  in
  Label
-> [Parameter]
-> CodeGen PTX ()
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh e))
forall aenv a.
Label
-> [Parameter]
-> CodeGen PTX ()
-> CodeGen PTX (IROpenAcc PTX aenv a)
makeOpenAcc Label
"stencil_border" ([Parameter]
paramFrom [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramInside [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramOut [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramIn [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramEnv) (CodeGen PTX () -> CodeGen PTX (IROpenAcc PTX aenv (Array sh e)))
-> CodeGen PTX () -> CodeGen PTX (IROpenAcc PTX aenv (Array sh e))
forall a b. (a -> b) -> a -> b
$ do

    Operands Int
start <- Operands Int -> CodeGen PTX (Operands Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Operands Int
liftInt Int
0)
    Operands Int
end   <- ShapeR sh -> Operands sh -> CodeGen PTX (Operands Int)
forall sh arch.
ShapeR sh -> Operands sh -> CodeGen arch (Operands Int)
shapeSize ShapeR sh
shr Operands sh
shInside

    Operands Int
-> Operands Int
-> (Operands Int -> CodeGen PTX ())
-> CodeGen PTX ()
imapFromTo Operands Int
start Operands Int
end ((Operands Int -> CodeGen PTX ()) -> CodeGen PTX ())
-> (Operands Int -> CodeGen PTX ()) -> CodeGen PTX ()
forall a b. (a -> b) -> a -> b
$ \Operands Int
i -> do

      Operands sh
ixIn  <- ShapeR sh
-> Operands sh -> Operands Int -> CodeGen PTX (Operands sh)
forall sh arch.
ShapeR sh
-> Operands sh -> Operands Int -> CodeGen arch (Operands sh)
indexOfInt ShapeR sh
shr Operands sh
shInside Operands Int
i    -- convert to multidimensional index of inside region
      Operands sh
ixOut <- ShapeR sh
-> Operands sh -> Operands sh -> CodeGen PTX (Operands sh)
forall sh.
ShapeR sh
-> Operands sh -> Operands sh -> CodeGen PTX (Operands sh)
offset ShapeR sh
shr Operands sh
ixIn Operands sh
shFrom       -- shift to multidimensional index of outside region
      Operands e
r     <- IRFun1 PTX aenv (sh -> e)
-> Operands sh -> IROpenExp PTX ((), sh) aenv e
forall arch env aenv a b.
IROpenFun1 arch env aenv (a -> b)
-> Operands a -> IROpenExp arch (env, a) aenv b
app1 IRFun1 PTX aenv (sh -> e)
apply Operands sh
Operands sh
ixOut             -- apply generator function
      Operands Int
j     <- ShapeR sh
-> Operands sh -> Operands sh -> CodeGen PTX (Operands Int)
forall sh arch.
ShapeR sh
-> Operands sh -> Operands sh -> CodeGen arch (Operands Int)
intOfIndex ShapeR sh
shr Operands sh
Operands sh
shOut Operands sh
ixOut
      IntegralType Int
-> IRArray (Array sh e)
-> Operands Int
-> Operands e
-> CodeGen PTX ()
forall int sh e arch.
IntegralType int
-> IRArray (Array sh e)
-> Operands int
-> Operands e
-> CodeGen arch ()
writeArray IntegralType Int
TypeInt IRArray (Array sh e)
arrOut Operands Int
j Operands e
r

    CodeGen PTX ()
forall arch. HasCallStack => CodeGen arch ()
return_


offset :: ShapeR sh -> Operands sh -> Operands sh -> CodeGen PTX (Operands sh)
offset :: ShapeR sh
-> Operands sh -> Operands sh -> CodeGen PTX (Operands sh)
offset ShapeR sh
shr Operands sh
sh1 Operands sh
sh2 = ShapeR sh
-> Operands sh -> Operands sh -> CodeGen PTX (Operands sh)
forall sh.
ShapeR sh
-> Operands sh -> Operands sh -> CodeGen PTX (Operands sh)
go ShapeR sh
shr Operands sh
sh1 Operands sh
sh2
  where
    go :: ShapeR t -> Operands t -> Operands t -> CodeGen PTX (Operands t)
    go :: ShapeR t -> Operands t -> Operands t -> CodeGen PTX (Operands t)
go ShapeR t
ShapeRz Operands t
OP_Unit Operands t
OP_Unit
      = Operands () -> CodeGen PTX (Operands ())
forall (m :: * -> *) a. Monad m => a -> m a
return Operands ()
OP_Unit

    go (ShapeRsnoc ShapeR sh1
t) (OP_Pair sa1 sb1) (OP_Pair sa2 sb2)
      = do Operands Int
x <- NumType Int
-> Operands Int -> Operands Int -> CodeGen PTX (Operands Int)
forall a arch.
NumType a -> Operands a -> Operands a -> CodeGen arch (Operands a)
add (NumType Int
forall a. IsNum a => NumType a
numType :: NumType Int) Operands Int
sb1 Operands Int
sb2
           Operands sh1 -> Operands Int -> Operands (sh1, Int)
forall a b. Operands a -> Operands b -> Operands (a, b)
OP_Pair (Operands sh1 -> Operands Int -> Operands (sh1, Int))
-> CodeGen PTX (Operands sh1)
-> CodeGen PTX (Operands Int -> Operands (sh1, Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ShapeR sh1
-> Operands sh1 -> Operands sh1 -> CodeGen PTX (Operands sh1)
forall sh.
ShapeR sh
-> Operands sh -> Operands sh -> CodeGen PTX (Operands sh)
go ShapeR sh1
t Operands sh1
sa1 Operands sh1
sa2 CodeGen PTX (Operands Int -> Operands (sh1, Int))
-> CodeGen PTX (Operands Int) -> CodeGen PTX (Operands (sh1, Int))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Operands Int -> CodeGen PTX (Operands Int)
forall (m :: * -> *) a. Monad m => a -> m a
return Operands Int
x