content: {-# LANGUAGE OverloadedStrings   #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  -- |
  -- Module      : Data.Array.Accelerate.LLVM.PTX.State
  -- 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.State (
  
    evalPTX,
    createTargetForDevice, createTargetFromContext,
  
    Pool(..),
    withPool, unsafeWithPool,
    defaultTarget,
    defaultTargetPool,
  
  ) where
  
  import Data.Array.Accelerate.Error
  
  import Data.Array.Accelerate.LLVM.State
  import Data.Array.Accelerate.LLVM.PTX.Target
  import qualified Data.Array.Accelerate.LLVM.PTX.Array.Table         as MT
  import qualified Data.Array.Accelerate.LLVM.PTX.Context             as CT
  import qualified Data.Array.Accelerate.LLVM.PTX.Debug               as Debug
  import qualified Data.Array.Accelerate.LLVM.PTX.Execute.Stream      as ST
  import qualified Data.Array.Accelerate.LLVM.PTX.Link.Cache          as LC
  import qualified Data.Array.Accelerate.LLVM.PTX.Pool                as Pool
  
  import Foreign.CUDA.Driver.Error
  import qualified Foreign.CUDA.Driver                                as CUDA
  import qualified Foreign.CUDA.Driver.Context                        as Context
  
  import Control.Exception                                            ( try, catch )
  import Data.Maybe                                                   ( fromMaybe, catMaybes )
  import Formatting
  import System.Environment                                           ( lookupEnv )
  import System.IO.Unsafe                                             ( unsafePerformIO, unsafeInterleaveIO )
  import Text.Read                                                    ( readMaybe )
  
  
  -- | Execute a PTX computation
  --
  evalPTX :: PTX -> LLVM PTX a -> IO a
  evalPTX ptx acc =
    CT.withContext (ptxContext ptx) (evalLLVM ptx acc)
    `catch`
    \e -> internalError shown (e :: CUDAException)
  
  
  -- | Create a new PTX execution target for the given device
  --
  createTargetForDevice
      :: CUDA.Device
      -> CUDA.DeviceProperties
      -> [CUDA.ContextFlag]
      -> IO PTX
  createTargetForDevice dev prp flags = do
    raw <- CUDA.create dev flags
    ptx <- createTarget dev prp raw
    _   <- CUDA.pop
    return ptx
  
  
  -- | Create a PTX execute target for the given device context
  --
  createTargetFromContext
      :: CUDA.Context
      -> IO PTX
  createTargetFromContext raw = do
    dev <- Context.device
    prp <- CUDA.props dev
    createTarget dev prp raw
  
  
  -- | Create a PTX execution target
  --
  createTarget
      :: CUDA.Device
      -> CUDA.DeviceProperties
      -> CUDA.Context
      -> IO PTX
  createTarget dev prp raw = do
    ctx <- CT.raw dev prp raw
    mt  <- MT.new ctx
    lc  <- LC.new
    st  <- ST.new ctx
    return $! PTX ctx mt lc st
  
  
  -- Shared execution contexts
  -- -------------------------
  
  -- In order to implement runN, we need to keep track of all available contexts,
  -- as well as the managed resource pool.
  --
  data Pool a = Pool
      { managed   :: {-# UNPACK #-} !(Pool.Pool a)
      , unmanaged :: [a]
      }
  
  -- Evaluate a thing given an execution context from the default pool
  --
  withPool :: Pool a -> (a -> IO b) -> IO b
  withPool p = Pool.with (managed p)
  
  unsafeWithPool :: Pool a -> (a -> b) -> b
  unsafeWithPool p = Pool.unsafeWith (managed p)
  
  
  -- 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.
  --
  
  -- | Select a device from the default pool.
  --
  {-# NOINLINE defaultTarget #-}
  defaultTarget :: PTX
  defaultTarget = head (unmanaged defaultTargetPool)
  
  -- | Create a shared resource pool of the available CUDA devices.
  --
  -- This globally shared resource pool is auto-initialised on startup. It will
  -- consist of every currently available device, or those specified by the value
  -- of the environment variable @ACCELERATE_LLVM_PTX_DEVICES@ (as a list of
  -- device ordinals).
  --
  {-# NOINLINE defaultTargetPool #-}
  defaultTargetPool :: Pool PTX
  defaultTargetPool = unsafePerformIO $! do
    Debug.traceM Debug.dump_gc "gc: initialise default PTX pool"
    CUDA.initialise []
  
    -- Figure out which GPUs we should put into the execution pool
    --
    ngpu  <- CUDA.count
    menv  <- (readMaybe =<<) <$> lookupEnv "ACCELERATE_LLVM_PTX_DEVICES"
  
    let ids = fromMaybe [0..ngpu-1] menv
  
        -- Spin up the GPU at the given ordinal.
        --
        boot :: Int -> IO (Maybe PTX)
        boot i = unsafeInterleaveIO $ do
          dev <- CUDA.device i
          prp <- CUDA.props dev
          r   <- try $ createTargetForDevice dev prp [CUDA.SchedAuto]
          case r of
            Right ptx               -> return (Just ptx)
            Left (e::CUDAException) -> do
              Debug.traceM Debug.dump_gc ("gc: failed to initialise device " % int % ": " % shown) i e
              return Nothing
  
    -- Create the pool from the available devices, which get spun-up lazily as
    -- required (due to the implementation of the Pool, we will look ahead by one
    -- each time one device is requested).
    --
    devices <- catMaybes <$> mapM boot ids
    if null devices
      then error "No CUDA-capable devices are available"
      else Pool <$> Pool.create devices
                <*> return devices
  
  