{-# LANGUAGE BangPatterns      #-}
{-# LANGUAGE CPP               #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ViewPatterns      #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- |
-- Module      : Data.Array.Accelerate.CUDA.Execute.Stream
-- Copyright   : [2013..2014] Manuel M T Chakravarty, Gabriele Keller, Trevor L. McDonell
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <tmcdonell@cse.unsw.edu.au>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--

module Data.Array.Accelerate.CUDA.Execute.Stream (

  Stream, Reservoir, new, streaming,

) where

-- friends
import Data.Array.Accelerate.CUDA.Context                       ( Context(..) )
import Data.Array.Accelerate.CUDA.Execute.Event                 ( Event, EventTable )
import Data.Array.Accelerate.FullList                           ( FullList(..) )
import Data.Array.Accelerate.Lifetime                           ( Lifetime, withLifetime )
import qualified Data.Array.Accelerate.CUDA.Execute.Event       as Event
import qualified Data.Array.Accelerate.CUDA.Debug               as D
import qualified Data.Array.Accelerate.FullList                 as FL

-- libraries
import Control.Monad.Trans                                      ( MonadIO, liftIO )
import Control.Exception                                        ( bracket_ )
import Control.Concurrent.MVar                                  ( MVar, newMVar, withMVar, mkWeakMVar )
import System.Mem.Weak                                          ( Weak, deRefWeak )
import Foreign.CUDA.Driver.Stream                               ( Stream(..) )
import qualified Foreign.CUDA.Driver                            as CUDA
import qualified Foreign.CUDA.Driver.Stream                     as Stream

import qualified Data.HashTable.IO                              as HT


-- Representation
-- --------------

-- The Reservoir is a place to store CUDA execution streams that are currently
-- inactive. When a new stream is requested one is provided from the reservoir
-- if available, otherwise a fresh execution stream is created.
--
type HashTable key val  = HT.BasicHashTable key val

type RSV                = MVar ( HashTable (Lifetime CUDA.Context) (FullList () Stream) )
data Reservoir          = Reservoir {-# UNPACK #-} !RSV
                                    {-# UNPACK #-} !(Weak RSV)


-- Executing operations in streams
-- -------------------------------

-- Execute an operation in a unique execution stream. The (asynchronous) result
-- is passed to a second operation together with an event that will be signalled
-- once the operation is complete. The stream and event are released after the
-- second operation completes.
--
{-# INLINE streaming #-}
streaming :: MonadIO m => Context -> Reservoir -> EventTable -> (Stream -> m a) -> (Event -> a -> m b) -> m b
streaming !ctx !rsv@(Reservoir !_ !weak_rsv) !etbl !action !after = do
  stream <- liftIO $ create ctx rsv
  first  <- action stream
  end    <- liftIO $ Event.waypoint ctx etbl stream
  final  <- after end first
  liftIO $! destroy (weakContext ctx) weak_rsv stream
  liftIO $! Event.destroy end
  return final


-- Primitive operations
-- --------------------

-- Generate a new empty reservoir. It is not necessary to pre-populate it with
-- any streams because stream creation does not cause a device synchronisation.
--
new :: IO Reservoir
new = do
  tbl    <- HT.new
  ref    <- newMVar tbl
  weak   <- mkWeakMVar ref (flush tbl)
  return $! Reservoir ref weak


-- Create a CUDA execution stream. If an inactive stream is available for use,
-- that is returned, else a fresh stream is created.
--
{-# INLINE create #-}
create :: Context -> Reservoir -> IO Stream
create !ctx (Reservoir !ref !_) = withMVar ref $ \tbl -> do
  --
  let key = deviceContext ctx
  ms    <- HT.lookup tbl key
  case ms of
    Nothing -> do
      stream <- Stream.create []
      message ("new " ++ showStream stream)
      return stream

    Just (FL () stream rest) -> do
      case rest of
        FL.Nil          -> HT.delete tbl key
        FL.Cons () s ss -> HT.insert tbl key (FL () s ss)
        --
      return stream


-- Merge a stream back into the reservoir. This must only be done once all
-- pending operations in the stream have completed.
--
{-# INLINE destroy #-}
destroy :: Weak (Lifetime CUDA.Context) -> Weak RSV -> Stream -> IO ()
destroy !weak_ctx !weak_rsv !stream = do
  -- Wait for all preceding operations submitted to the stream to complete. Not
  -- necessary because of the setup of 'streaming'.
  -- Stream.block stream

  -- Now check whether the context and reservoir are still active. Return
  -- the stream back to the reservoir for later reuse if we can, otherwise
  -- destroy it.
  mc <- deRefWeak weak_ctx
  case mc of
    Nothing       -> message ("finalise/dead context " ++ showStream stream)
    Just ctx      -> do
      --
      mr <- deRefWeak weak_rsv
      case mr of
        Nothing   -> trace ("destroy/free " ++ showStream stream) $ Stream.destroy stream
        Just ref  -> trace ("destroy/save " ++ showStream stream) $ withMVar ref $ \tbl -> do
          --
          ms <- HT.lookup tbl ctx
          case ms of
            Nothing       -> HT.insert tbl ctx (FL.singleton () stream)
            Just ss       -> HT.insert tbl ctx (FL.cons () stream ss)


-- Add a finaliser to an execution stream
--
-- addStreamFinalizer :: Stream -> IO () -> IO ()
-- addStreamFinalizer st@(Stream (Ptr st#)) f = IO $ \s ->
--   case mkWeak# st# st f s of (# s', _w #) -> (# s', () #)


-- Destroy all streams in the reservoir.
--
flush :: HashTable (Lifetime CUDA.Context) (FullList () Stream) -> IO ()
flush !tbl =
  let clean (!lctx,!ss) = do
        withLifetime lctx $ \ctx -> bracket_ (CUDA.push ctx) CUDA.pop (FL.mapM_ (const Stream.destroy) ss)
        HT.delete tbl lctx
  in
  message "flush reservoir" >> HT.mapM_ clean tbl


-- Debug
-- -----

{-# INLINE trace #-}
trace :: String -> IO a -> IO a
trace msg next = message msg >> next

{-# INLINE message #-}
message :: String -> IO ()
message msg = D.traceIO D.dump_sched ("stream: " ++ msg)

{-# INLINE showStream #-}
showStream :: Stream -> String
showStream (Stream s) = show s