content: {-# 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.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 )
  
  
  -- Apply a unary function to each element of an array. Each thread processes
  -- multiple elements, striding the array by the grid size.
  --
  mkMap :: UID
        -> Gamma       aenv
        -> ArrayR (Array sh a)
        -> TypeR b
        -> IRFun1  PTX aenv (a -> b)
        -> CodeGen PTX      (IROpenAcc PTX aenv (Array sh b))
  mkMap uid aenv repr@(ArrayR shr _) tp' apply =
    let
        (arrOut, paramOut)  = mutableArray (ArrayR shr tp') "out"
        (arrIn,  paramIn)   = mutableArray repr             "in"
        paramEnv            = envParam aenv
    in
    makeOpenAcc uid "map" (paramOut ++ paramIn ++ paramEnv) $ do
  
      start <- return (liftInt 0)
      end   <- shapeSize shr (irArrayShape arrIn)
  
      imapFromTo start end $ \i -> do
        xs <- readArray TypeInt arrIn i
        ys <- app1 apply xs
        writeArray TypeInt arrOut i ys
  
      return_
  
  