content: {-# LANGUAGE BangPatterns        #-}
  {-# LANGUAGE ConstraintKinds     #-}
  {-# LANGUAGE FlexibleInstances   #-}
  {-# LANGUAGE GADTs               #-}
  {-# LANGUAGE InstanceSigs        #-}
  {-# LANGUAGE PatternGuards       #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TemplateHaskell     #-}
  {-# LANGUAGE TypeFamilies        #-}
  {-# OPTIONS_GHC -fno-warn-orphans #-}
  -- |
  -- Module      : Data.Array.Accelerate.CUDA.Array.Remote
  -- Copyright   : [2015..2016] Manuel M T Chakravarty, Gabriele Keller, Trevor L. McDonell, Robert Clifton-Everest
  -- License     : BSD3
  --
  -- Maintainer  : Robert Clifton-Everest <robertce@cse.unsw.edu.au>
  -- Stability   : experimental
  -- Portability : non-portable (GHC extensions)
  --
  
  module Data.Array.Accelerate.CUDA.Array.Remote (
  
    -- Tables for host/device memory associations
    MemoryTable, R.PrimElt,
    new, malloc, withRemote, free, insertUnmanaged, reclaim
  
  ) where
  
  import Control.Concurrent.MVar                                  ( MVar, newMVar, withMVar, modifyMVar, readMVar )
  import Control.Exception
  import Control.Monad.IO.Class                                   ( MonadIO, liftIO )
  import Control.Monad.Trans.Reader
  import Data.Functor
  import Data.IntMap.Strict                                       ( IntMap )
  import Data.Proxy
  import Data.Typeable                                            ( Typeable )
  import Foreign.Ptr                                              ( ptrToIntPtr )
  import Foreign.Storable
  import Prelude                                                  hiding ( lookup )
  import qualified Data.IntMap.Strict                             as IM
  
  import Foreign.CUDA.Driver.Error
  import qualified Foreign.CUDA.Driver                            as CUDA
  
  import Data.Array.Accelerate.Array.Data
  import Data.Array.Accelerate.Error
  import Data.Array.Accelerate.Lifetime                           ( unsafeGetValue )
  import qualified Data.Array.Accelerate.Array.Remote             as R
  
  import Data.Array.Accelerate.CUDA.Context                       ( Context(..), push, pop )
  import Data.Array.Accelerate.CUDA.Execute.Event                 ( Event, EventTable, waypoint, query )
  import Data.Array.Accelerate.CUDA.Execute.Stream                ( Stream )
  import qualified Data.Array.Accelerate.CUDA.Debug               as Debug
  
  
  -- Instance for basic remote memory management functionality.
  --
  type CRM = ReaderT (Maybe Stream) IO
  
  instance R.RemoteMemory CRM where
    type RemotePtr CRM = CUDA.DevicePtr
    --
    mallocRemote n       = ReaderT $ \_ ->
      fmap Just (CUDA.mallocArray n)
        `catch` \e -> case e of
                        ExitCode OutOfMemory -> return Nothing
                        _                    -> trace ("malloc failed with error: " ++ show e) (throwIO e)
  
    pokeRemote n dst ad  = ReaderT $ \mst ->
      transfer "poke" (n * sizeOfPtr dst) $
      CUDA.pokeArrayAsync n (CUDA.HostPtr (ptrsOfArrayData ad)) dst mst
  
    peekRemote n src ad  = ReaderT $ \mst ->
      transfer "peek" (n * sizeOfPtr src) $
      CUDA.peekArrayAsync n src (CUDA.HostPtr (ptrsOfArrayData ad)) mst
  
    castRemotePtr _      = CUDA.castDevPtr
    totalRemoteMem       = ReaderT $ \_ -> snd <$> CUDA.getMemInfo
    availableRemoteMem   = ReaderT $ \_ -> fst <$> CUDA.getMemInfo
    remoteAllocationSize = return 1024
  
  
  -- We leverage the memory cache from the accelerate base package.
  --
  -- However, we actually need multiple caches because every pointer has an
  -- associated CUDA context. We could pair every DevicePtr with its context and
  -- just have a single table, but the LRU implementation in the base package
  -- assumes that remote pointers can be re-used, something that would not be true
  -- for pointers allocated under different contexts.
  --
  type MT          = IntMap (R.MemoryTable CUDA.DevicePtr Task)
  data MemoryTable = MemoryTable {-# UNPACK #-} !EventTable
                                 {-# UNPACK #-} !(MVar MT)
  
  type Task = Maybe Event
  
  instance R.Task Task where
    completed Nothing  = return True
    completed (Just e) = query e
  
  
  -- Create a MemoryTable.
  --
  new :: EventTable -> IO MemoryTable
  new et = do
    message "initialise CUDA memory table"
    MemoryTable et <$> newMVar IM.empty
  
  
  -- Perform action on the device ptr that matches the given host-side array. Any
  -- operations
  --
  withRemote
      :: forall e a b. R.PrimElt e a
      => Context
      -> MemoryTable
      -> ArrayData e
      -> (CUDA.DevicePtr a -> IO b)
      -> Maybe Stream
      -> IO (Maybe b)
  withRemote ctx (MemoryTable et ref) ad run ms = do
    ct <- readMVar ref
    case IM.lookup (contextId ctx) ct of
      Nothing -> $internalError "withRemote" "context not found"
      Just mc -> streaming ms $ R.withRemote mc ad run'
    where
      run' :: R.RemotePtr CRM a -> CRM (Task, b)
      run' p = liftIO $ do
        c  <- run p
        case ms of
          Nothing -> return (Nothing, c)
          Just s  -> do
            e <- waypoint ctx et s
            return (Just e, c)
  
  
  -- Allocate a new device array to be associated with the given host-side array.
  -- Has the same properties as `Data.Array.Accelerate.Array.Remote.LRU.malloc`
  malloc :: forall a b. (Typeable a, R.PrimElt a b)
         => Context
         -> MemoryTable
         -> ArrayData a
         -> Bool
         -> Int
         -> IO Bool
  malloc !ctx (MemoryTable _ !ref) !ad !frozen !n = do
    mt <- modifyMVar ref $ \ct -> blocking $ do
     case IM.lookup (contextId ctx) ct of
             Nothing -> trace "malloc/context not found" $ insertContext ctx ct
             Just mt -> return (ct, mt)
    blocking $ R.malloc mt ad frozen n
  
  
  -- Explicitly free an array in the LRU table. Has the same properties as
  -- `Data.Array.Accelerate.Array.Remote.LRU.free`
  --
  free :: R.PrimElt a b
       => Context
       -> MemoryTable
       -> ArrayData a
       -> IO ()
  free !ctx (MemoryTable _ !ref) !arr = withMVar ref $ \ct ->
    case IM.lookup (contextId ctx) ct of
      Nothing -> message "free/context not found"
      Just mt -> R.free (Proxy :: Proxy CRM) mt arr
  
  
  -- Record an association between a host-side array and a device memory area that was
  -- not allocated by accelerate. The device memory will NOT be freed by the memory
  -- manager.
  --
  insertUnmanaged
      :: R.PrimElt a b
      => Context
      -> MemoryTable
      -> ArrayData a
      -> CUDA.DevicePtr b
      -> IO ()
  insertUnmanaged !ctx (MemoryTable _ !ref) !arr !ptr = do
    mt <- modifyMVar ref $ \ct -> blocking $ do
     case IM.lookup (contextId ctx) ct of
             Nothing -> trace "insertUnmanaged/context not found" $ insertContext ctx ct
             Just mt -> return (ct, mt)
    blocking $ R.insertUnmanaged mt arr ptr
  
  insertContext
      :: Context
      -> MT
      -> CRM ( MT, R.MemoryTable CUDA.DevicePtr Task )
  insertContext ctx ct = liftIO $ do
     mt <- R.new (\p -> bracket_ (push ctx) pop (CUDA.free p))
     return (IM.insert (contextId ctx) mt ct, mt)
  
  
  -- Removing entries
  -- ----------------
  
  -- Initiate garbage collection and finalise any arrays that have been marked as
  -- unreachable.
  --
  reclaim :: MemoryTable -> IO ()
  reclaim (MemoryTable _ ref) = withMVar ref (blocking . mapM_ R.reclaim . IM.elems)
  
  -- Miscellaneous
  -- -------------
  
  {-# INLINE contextId #-}
  contextId :: Context -> Int
  contextId !ctx =
    let CUDA.Context !p = unsafeGetValue (deviceContext ctx)
    in fromIntegral (ptrToIntPtr p)
  
  {-# INLINE blocking #-}
  blocking :: CRM a -> IO a
  blocking = flip runReaderT Nothing
  
  {-# INLINE streaming #-}
  streaming :: Maybe Stream -> CRM a -> IO a
  streaming = flip runReaderT
  
  {-# INLINE sizeOfPtr #-}
  sizeOfPtr :: forall a. Storable a => CUDA.DevicePtr a -> Int
  sizeOfPtr _ = sizeOf (undefined :: a)
  
  -- Debug
  -- -----
  
  {-# INLINE trace #-}
  trace :: MonadIO m => String -> m a -> m a
  trace msg next = message msg >> next
  
  {-# INLINE message #-}
  message :: MonadIO m => String -> m ()
  message msg = liftIO $ Debug.traceIO Debug.dump_gc ("gc/lru: " ++ msg)
  
  {-# INLINE showBytes #-}
  showBytes :: Int -> String
  showBytes x = Debug.showFFloatSIBase (Just 0) 1024 (fromIntegral x :: Double) "B"
  
  {-# INLINE transfer #-}
  transfer :: String -> Int -> IO () -> IO ()
  transfer name bytes action
    = let showRate x t        = Debug.showFFloatSIBase (Just 3) 1024 (fromIntegral x / t) "B/s"
          msg gpuTime cpuTime = "gc/lru: " ++ name ++ ": "
                                           ++ showBytes bytes ++ " @ " ++ showRate bytes gpuTime ++ ", "
                                           ++ Debug.elapsed gpuTime cpuTime
      in
      Debug.timed Debug.dump_gc msg Nothing action
  
  