content: {-# LANGUAGE BangPatterns               #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE ScopedTypeVariables        #-}
  {-# LANGUAGE TemplateHaskell            #-}
  -- |
  -- Module      : Data.Array.Accelerate.CUDA.State
  -- Copyright   : [2008..2014] Manuel M T Chakravarty, Gabriele Keller
  --               [2009..2014] Trevor L. McDonell
  -- License     : BSD3
  --
  -- Maintainer  : Trevor L. McDonell <tmcdonell@cse.unsw.edu.au>
  -- Stability   : experimental
  -- Portability : non-portable (GHC extensions)
  --
  -- This module defines a state monad token which keeps track of the code
  -- generator state, including memory transfers and external compilation
  -- processes.
  --
  
  module Data.Array.Accelerate.CUDA.State (
  
    -- Evaluating computations
    CIO, Context, evalCUDA, evalCUDAState,
  
    -- Querying execution state
    defaultContext, deviceProperties, activeContext, kernelTable, memoryTable, streamReservoir,
    eventTable,
  
  ) where
  
  -- friends
  import Data.Array.Accelerate.Error
  import Data.Array.Accelerate.CUDA.Context
  import Data.Array.Accelerate.CUDA.Debug                 ( traceIO, dump_gc )
  import Data.Array.Accelerate.CUDA.Persistent            as KT ( KernelTable, new )
  import Data.Array.Accelerate.CUDA.Array.Remote          as MT ( MemoryTable, new )
  import Data.Array.Accelerate.CUDA.Execute.Stream        as ST ( Reservoir, new )
  import Data.Array.Accelerate.CUDA.Execute.Event         as ET ( EventTable, new)
  import Data.Array.Accelerate.CUDA.Analysis.Device
  
  -- library
  import Control.Applicative
  import Control.Concurrent                               ( runInBoundThread )
  import Control.Exception                                ( catch, bracket_ )
  import Control.Monad.Reader                             ( MonadReader, ReaderT(..), runReaderT )
  import Control.Monad.State.Strict                       ( MonadState, StateT(..), evalStateT )
  import Control.Monad.Trans                              ( MonadIO )
  import System.IO.Unsafe                                 ( unsafePerformIO )
  import Foreign.CUDA.Driver.Error
  import qualified Foreign.CUDA.Driver                    as CUDA
  import Prelude
  
  
  -- Execution State
  -- ---------------
  
  -- The state token for CUDA accelerated array operations. This is a stack of
  -- (read only) device properties and context, and mutable state for tracking
  -- device memory and kernel object code.
  --
  data State = State {
      memoryTable         :: {-# UNPACK #-} !MemoryTable,                 -- host/device memory associations
      kernelTable         :: {-# UNPACK #-} !KernelTable,                 -- compiled kernel object code
      streamReservoir     :: {-# UNPACK #-} !Reservoir,                   -- kernel execution streams
      eventTable          :: {-# UNPACK #-} !EventTable                   -- CUDA events
    }
  
  newtype CIO a = CIO {
      runCIO              :: ReaderT Context (StateT State IO) a
    }
    deriving ( Functor, Applicative, Monad, MonadIO
             , MonadReader Context, MonadState State )
  
  
  -- Extract the active context from the execution state
  --
  {-# INLINE activeContext #-}
  activeContext :: Context -> Context
  activeContext = id
  
  -- |Evaluate a CUDA array computation
  --
  {-# NOINLINE evalCUDA #-}
  evalCUDA :: Context -> CIO a -> IO a
  evalCUDA !ctx !acc =
    runInBoundThread (bracket_ setup teardown action)
    `catch`
    \e -> $internalError "unhandled" (show (e :: CUDAException))
    where
      setup       = push ctx
      teardown    = pop
      action      = evalStateT (runReaderT (runCIO acc) ctx) theState
  
  -- |Evaluate a CUDA array computation with the specific state. Exceptions are
  -- not caught.
  --
  -- RCE: This is unfortunately hacky, but necessary to stop device pointers
  -- leaking.
  evalCUDAState :: Context -> MemoryTable -> KernelTable -> Reservoir -> EventTable -> CIO a -> IO a
  evalCUDAState ctx mt kt rsv etbl acc = evalStateT (runReaderT (runCIO acc) ctx)
                                                    (State mt kt rsv etbl)
  
  -- Top-level mutable state
  -- -----------------------
  --
  -- It is important to keep some information alive for the entire run of the
  -- program, not just a single execution. These tokens use unsafePerformIO to
  -- ensure they are executed only once, and reused for subsequent invocations.
  --
  {-# NOINLINE theState #-}
  theState :: State
  theState
    = unsafePerformIO
    $ do  message "initialise CUDA state"
          etbl    <- keepAlive =<< ET.new
          mtb     <- keepAlive =<< MT.new etbl
          ktb     <- keepAlive =<< KT.new
          rsv     <- keepAlive =<< ST.new
          return  $! State mtb ktb rsv etbl
  
  
  -- Select and initialise a default CUDA device, and create a new execution
  -- context. The device is selected based on compute capability and estimated
  -- maximum throughput.
  --
  {-# NOINLINE defaultContext #-}
  defaultContext :: Context
  defaultContext = unsafePerformIO $ do
    message "initialise default context"
    CUDA.initialise []
    (dev,_)       <- selectBestDevice
    create dev [CUDA.SchedAuto]
  
  
  -- Debug
  -- -----
  
  {-# INLINE message #-}
  message :: String -> IO ()
  message msg = traceIO dump_gc ("gc: " ++ msg)
  
  