{-# LANGUAGE GADTs               #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
-- |
-- Module      : Data.Array.Accelerate.LLVM.PTX.CodeGen.Map
-- Copyright   : [2014..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.Map
  where

-- accelerate
import Data.Array.Accelerate.Representation.Array
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.Monad
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 )


-- Apply a unary function to each element of an array. Each thread processes
-- multiple elements, striding the array by the grid size.
--
mkMap :: Gamma       aenv
      -> ArrayR (Array sh a)
      -> TypeR b
      -> IRFun1  PTX aenv (a -> b)
      -> CodeGen PTX      (IROpenAcc PTX aenv (Array sh b))
mkMap :: Gamma aenv
-> ArrayR (Array sh a)
-> TypeR b
-> IRFun1 PTX aenv (a -> b)
-> CodeGen PTX (IROpenAcc PTX aenv (Array sh b))
mkMap Gamma aenv
aenv repr :: ArrayR (Array sh a)
repr@(ArrayR ShapeR sh
shr TypeR e
_) TypeR b
tp' IRFun1 PTX aenv (a -> b)
apply =
  let
      (IRArray (Array sh b)
arrOut, [Parameter]
paramOut)  = ArrayR (Array sh b)
-> Name (Array sh b) -> (IRArray (Array sh b), [Parameter])
forall sh e.
ArrayR (Array sh e)
-> Name (Array sh e) -> (IRArray (Array sh e), [Parameter])
mutableArray (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') Name (Array sh b)
"out"
      (IRArray (Array sh a)
arrIn,  [Parameter]
paramIn)   = ArrayR (Array sh a)
-> Name (Array sh a) -> (IRArray (Array sh a), [Parameter])
forall sh e.
ArrayR (Array sh e)
-> Name (Array sh e) -> (IRArray (Array sh e), [Parameter])
mutableArray ArrayR (Array sh a)
repr             Name (Array sh a)
"in"
      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 b))
forall aenv a.
Label
-> [Parameter]
-> CodeGen PTX ()
-> CodeGen PTX (IROpenAcc PTX aenv a)
makeOpenAcc Label
"map" ([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 b)))
-> CodeGen PTX () -> CodeGen PTX (IROpenAcc PTX aenv (Array sh b))
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 (IRArray (Array sh a) -> Operands sh
forall sh e. IRArray (Array sh e) -> Operands sh
irArrayShape IRArray (Array sh a)
arrIn)

    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 a
xs <- IntegralType Int
-> IRArray (Array sh a) -> Operands Int -> CodeGen PTX (Operands a)
forall int sh e arch.
IntegralType int
-> IRArray (Array sh e)
-> Operands int
-> CodeGen arch (Operands e)
readArray IntegralType Int
TypeInt IRArray (Array sh a)
arrIn Operands Int
i
      Operands b
ys <- IRFun1 PTX aenv (a -> b)
-> Operands a -> IROpenExp PTX ((), a) 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 (a -> b)
apply Operands a
xs
      IntegralType Int
-> IRArray (Array sh b)
-> Operands Int
-> Operands b
-> 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 b)
arrOut Operands Int
i Operands b
ys

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