content: {-# LANGUAGE BangPatterns        #-}
  {-# LANGUAGE GADTs               #-}
  {-# LANGUAGE MagicHash           #-}
  {-# LANGUAGE OverloadedStrings   #-}
  {-# LANGUAGE RankNTypes          #-}
  {-# LANGUAGE RecordWildCards     #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TemplateHaskell     #-}
  {-# LANGUAGE UnboxedTuples       #-}
  {-# OPTIONS_GHC -Wno-orphans #-}
  -- |
  -- Module      : Data.Array.Accelerate.LLVM.PTX.Array.Data
  -- 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.Array.Data (
  
    module Data.Array.Accelerate.LLVM.Array.Data,
    module Data.Array.Accelerate.LLVM.PTX.Array.Data,
  
  ) where
  
  import Data.Array.Accelerate.Array.Data
  import Data.Array.Accelerate.Array.Unique
  import Data.Array.Accelerate.Error
  import Data.Array.Accelerate.Lifetime
  import Data.Array.Accelerate.Representation.Array
  import Data.Array.Accelerate.Representation.Shape
  import Data.Array.Accelerate.Representation.Type
  import Data.Array.Accelerate.Type
  
  import Data.Array.Accelerate.LLVM.Array.Data
  import Data.Array.Accelerate.LLVM.State
  
  import Data.Array.Accelerate.LLVM.PTX.State
  import Data.Array.Accelerate.LLVM.PTX.Target
  import Data.Array.Accelerate.LLVM.PTX.Execute.Async
  import qualified Data.Array.Accelerate.LLVM.PTX.Array.Prim          as Prim
  
  import Control.Applicative
  import Control.Monad
  import Control.Monad.Reader
  import Control.Monad.State                                          ( gets )
  import System.IO.Unsafe
  import Prelude
  
  import GHC.Heap.NormalForm
  
  
  -- | Remote memory management for the PTX target. Data can be copied
  -- asynchronously using multiple execution engines whenever possible.
  --
  instance Remote PTX where
    {-# INLINEABLE allocateRemote   #-}
    {-# INLINEABLE indexRemoteAsync #-}
    {-# INLINEABLE useRemoteR       #-}
    {-# INLINEABLE copyToHostR      #-}
    {-# INLINEABLE copyToRemoteR    #-}
    allocateRemote repr@(ArrayR shr tp) !sh = do
      let !n = size shr sh
      arr    <- liftIO $ allocateArray repr sh  -- shadow array on the host
      liftPar $ runArray tp arr (\m t ad -> Prim.mallocArray t (n*m) ad >> return ad)
  
    indexRemoteAsync  = runIndexArrayAsync Prim.indexArrayAsync
    useRemoteR        = Prim.useArrayAsync
    copyToHostR       = Prim.peekArrayAsync
    copyToRemoteR     = Prim.pokeArrayAsync
    copyToPeerR       = internalError "not supported yet"
  
  
  -- | Copy an array from the remote device to the host. Although the Accelerate
  -- program is hyper-strict and will evaluate the computation as soon as any part
  -- of it is demanded, the individual array payloads are copied back to the host
  -- _only_ as they are demanded by the Haskell program. This has several
  -- consequences:
  --
  --   1. If the device has multiple memcpy engines, only one will be used. The
  --      transfers are however associated with a non-default stream.
  --
  --   2. Using 'seq' to force an Array to head-normal form will initiate the
  --      computation, but not transfer the results back to the host. Requesting
  --      an array element or using 'deepseq' to force to normal form is required
  --      to actually transfer the data.
  --
  {-# INLINEABLE copyToHostLazy #-}
  copyToHostLazy
      :: HasCallStack
      => ArraysR arrs
      -> FutureArraysR PTX arrs
      -> Par PTX arrs
  copyToHostLazy TupRunit         ()       = return ()
  copyToHostLazy (TupRpair r1 r2) (f1, f2) = do
    a1 <- copyToHostLazy r1 f1
    a2 <- copyToHostLazy r2 f2
    return (a1, a2)
  copyToHostLazy (TupRsingle (ArrayR shr tp)) future = do
    ptx <- gets llvmTarget
    liftIO $ do
      Array sh adata <- wait future
  
      -- Note: [Lazy device-host transfers]
      --
      -- This needs must be non-strict at the leaves of the datatype (that
      -- is, the UniqueArray pointers). This means we can traverse the
      -- ArrayData constructors (in particular, the spine defined by Unit
      -- and Pair) until we reach the array we care about, without forcing
      -- the other fields.
      --
      -- https://github.com/AccelerateHS/accelerate/issues/437
      --
      -- Furthermore, we only want to transfer the data if the host pointer
      -- is currently unevaluated. This situation can occur for example if
      -- the argument to 'use' or 'unit' is returned as part of the result
      -- of a 'run'. Peek at GHC's underlying closure representation and
      -- check whether the pointer is a thunk, and only initiate the
      -- transfer if so.
      --
      let
        peekR :: SingleType e
              -> ArrayData e
              -> Int
              -> IO (ArrayData e)
        peekR t ad m
          | SingleArrayDict                        <- singleArrayDict t
          , UniqueArray uid (Lifetime lft weak fp) <- ad
          = unsafeInterleaveIO $ do
            yes <- isNormalForm fp
            fp' <- if yes
                      then return fp
                      else unsafeInterleaveIO . evalPTX ptx . evalPar $ do
                            !_ <- block =<< Prim.peekArrayAsync t m ad
                            return fp
            --
            return $ UniqueArray uid (Lifetime lft weak fp')
  
        n = size shr sh
  
        runR :: TypeR e -> ArrayData e -> IO (ArrayData e)
        runR TupRunit           !()          = return ()
        runR (TupRpair !t1 !t2) (!ad1, !ad2) = (,) <$> runR t1 ad1 <*> runR t2 ad2
        runR (TupRsingle !t)    !ad          =
          case t of
            SingleScalarType s                       -> peekR s ad n
            VectorScalarType (VectorType w s)
              | SingleArrayDict <- singleArrayDict s -> peekR s ad (n * w)
  
      Array sh <$> runR tp adata
  
  -- | Clone an array into a newly allocated array on the device.
  --
  cloneArrayAsync
      :: ArrayR (Array sh e)
      -> Array sh e
      -> Par PTX (Future (Array sh e))
  cloneArrayAsync repr@(ArrayR shr tp) arr@(Array _ src) = do
    Array _ dst <- allocateRemote repr sh
    Array sh `liftF` copyR tp src dst
    where
      sh = shape arr
      n  = size shr sh
  
      copyR :: TypeR s -> ArrayData s -> ArrayData s -> Par PTX (Future (ArrayData s))
      copyR TupRunit           !_          !_            = newFull ()
      copyR (TupRpair !t1 !t2) !(ad1, ad2) !(ad1', ad2') = liftF2 (,) (copyR t1 ad1 ad1') (copyR t2 ad2 ad2')
      copyR (TupRsingle !t)    !ad         !ad'          =
        case t of
          SingleScalarType s                       -> copyPrim s ad ad' n
          VectorScalarType (VectorType w s)
            | SingleArrayDict <- singleArrayDict s -> copyPrim s ad ad' (n * w)
  
      copyPrim
          :: SingleType s
          -> ArrayData s
          -> ArrayData s
          -> Int
          -> Par PTX (Future (ArrayData s))
      copyPrim !s !a1 !a2 !m = Prim.copyArrayAsync s m a1 a2
  
      liftF :: Async arch
            => (a -> b)
            -> Par arch (FutureR arch a)
            -> Par arch (FutureR arch b)
      liftF f x = do
        r  <- new
        x' <- x
        put r . f =<< get x'  -- don't create a new execution stream for this
        return r
  
      liftF2 :: Async arch
             => (a -> b -> c)
             -> Par arch (FutureR arch a)
             -> Par arch (FutureR arch b)
             -> Par arch (FutureR arch c)
      liftF2 f x y = do
        r  <- new
        x' <- spawn x
        y' <- spawn y
        fork $ put r =<< liftM2 f (get x') (get y')
        return r
  
  