Context.hs

{-# LANGUAGE RankNTypes, ExistentialQuantification, FlexibleInstances, UndecidableInstances, TypeFamilyDependencies, MultiParamTypeClasses, ScopedTypeVariables, FlexibleContexts, AllowAmbiguousTypes #-}
{-|
Defines context options and utilities used internally.
Mostly useful for creating new abstractions from the @Raw@ interface. 
-}

module Futhask.Context where
import Data.Word
import Data.Int
import Foreign.C
import Foreign.C.String
import Foreign.Ptr
import Foreign.Storable
import Foreign.ForeignPtr (ForeignPtr, withForeignPtr)
import Foreign.Concurrent
import Foreign.Marshal.Alloc
import Foreign.Marshal.Array as A
import Control.Monad
import Control.Concurrent
import Control.Concurrent.MVar
import System.Mem

-- *Reference Counting
type ReferenceCounter = MVar Int

newRefCounter :: IO ReferenceCounter
newRefCounter = newMVar 0

increment :: ReferenceCounter -> IO ()
increment c = modifyMVar_ c (\a -> return $! (a+1))

decrement :: ReferenceCounter -> IO ()
decrement c = modifyMVar_ c (\a -> return $! (a-1))

isZero :: ReferenceCounter -> IO Bool
isZero c = readMVar c >>= \a -> return (a==0) 

waitForZero :: ReferenceCounter -> IO ()
waitForZero c = isZero c >>= \b -> if b then return () else yield >> waitForZero c
newForeignPtrWithRef refCounter pointer finalizer 
    = increment refCounter >> newForeignPtr pointer (finalizer >> decrement refCounter)
    
-- *Context

-- | Context options for the different backends.
-- Using an option in a backend that doesn't support it will simply do nothing.
data ContextOption
    = TuningParameters [(String, CSize)] -- ^All backends
    | CacheFile String -- ^All backends
    | Debug Int -- ^All backends
    | Log Int -- ^All backends
    | Profile Int -- ^All backends
    | Device String -- ^GPU backends
    | LoadProgram String -- ^GPU backends
    | DumpProgram String -- ^GPU backends
    | DefaultThreadBlockSize Int -- ^GPU backends
    | DefaultGridSize Int -- ^GPU backends
    | DefaultTileSize Int -- ^GPU backends
    | Platform String -- ^OpenCL
    | BuildOptions [String] -- ^OpenCL
    | NvrtcOptions [String] -- ^Cuda
    | NumThreads Int -- ^Multicore

-- |Each generated library has its own instance of @Context@ 
class Context c where
    type RawContext c = raw | raw -> c
    type RawConfig  c = cfg | cfg -> c
    -- wrapping
    wrapRawContext   :: (ReferenceCounter, ForeignPtr (RawContext c)) -> c
    unwrapRawContext :: c -> (ReferenceCounter, ForeignPtr (RawContext c))
    -- general
    newContext  :: Ptr (RawConfig c)  -> IO (Ptr (RawContext c))
    freeContext :: Ptr (RawContext c) -> IO ()
    clearCaches :: Ptr (RawContext c) -> IO Int 
    syncContext :: Ptr (RawContext c) -> IO Int
    -- reporting
    getError  :: Ptr (RawContext c) -> IO CString
    getReport :: Ptr (RawContext c) -> IO CString
    -- profiling 
    pauseProfiling   :: Ptr (RawContext c) -> IO ()
    unpauseProfiling :: Ptr (RawContext c) -> IO ()
    -- configuration
    newConfig  :: IO (Ptr (RawConfig c))
    freeConfig :: Ptr (RawConfig c) -> IO ()
    -- tuning
    setTuningParam      :: Ptr (RawConfig c) -> CString -> CSize -> IO Int
    -- general options
    setCacheFile :: Ptr (RawConfig c) -> CString -> IO ()
    setDebug     :: Ptr (RawConfig c) -> Int -> IO ()               
    setLog       :: Ptr (RawConfig c) -> Int -> IO ()
    setProfile   :: Ptr (RawConfig c) -> Int -> IO ()
    -- Cuda and OpenCL options
    setDevice                 :: Ptr (RawConfig c) -> CString -> IO ()
    setDefaultThreadBlockSize :: Ptr (RawConfig c) -> Int -> IO ()
    setDefaultGridSize    :: Ptr (RawConfig c) -> Int -> IO ()
    setDefaultTileSize    :: Ptr (RawConfig c) -> Int -> IO ()
    -- OpenCL options
    setPlatform    :: Ptr (RawConfig c) -> CString -> IO ()
    --setLoadBinary  :: Ptr (RawConfig c) -> CString -> IO ()
    --setDumpBinary  :: Ptr (RawConfig c) -> CString -> IO ()
    setBuildOption :: Ptr (RawConfig c) -> CString -> IO ()
    setNvrtcOption :: Ptr (RawConfig c) -> CString -> IO ()
    setNumThreads :: Ptr (RawConfig c) -> Int -> IO ()


-- |implementation of options not available for backend - does nothing
voidOption :: config -> parameter -> IO ()
voidOption _ _ = return ()

-- |sets option and returns any error messages and post config finalizer operations
setContextOption :: (Context c) => Ptr (RawConfig c) -> ContextOption -> IO (String, IO ())
setContextOption config option = case option of
    TuningParameters ps -> fmap (\(errs, fins) -> (concat errs, sequence_ fins))
        $ fmap unzip
        $ mapM (\(n, v) -> withCString n (\cn -> setTuningParam config cn v >>= \i -> if i == 0 
             then return ("", return ()) 
             else return ("Tuning parameter \"" ++ n ++ "\" could not be set.\n", return ()))) ps 
    CacheFile fn -> newCString fn >>= \cn -> setCacheFile config cn >> return ("", free cn)
    Debug i   -> nofin $ setDebug   config i
    Log i     -> nofin $ setLog     config i
    Profile i -> nofin $ setProfile config i
    Device s -> nofin $ withCString s (setDevice config)
    DefaultThreadBlockSize s -> nofin $ setDefaultThreadBlockSize config s
    DefaultGridSize s        -> nofin $ setDefaultGridSize        config s
    DefaultTileSize s        -> nofin $ setDefaultTileSize        config s
    Platform p    -> nofin $ withCString p  (setPlatform   config)
    BuildOptions os -> nofin $ mapM_ (\o -> withCString o (setBuildOption config)) os
    NvrtcOptions os -> nofin $ mapM_ (\o -> withCString o (setNvrtcOption config)) os
    NumThreads n -> nofin $ setNumThreads config n
    where nofin a = a >> return ("", return ())

-- | Make new context with a list of options
mkContext :: (Context context) => [ContextOption] -> IO (context, [String])
mkContext options = do
    config <- newConfig
    (messages, optionFinalizers) <- unzip <$> mapM (setContextOption config) options
    context <- newContext config
    refCounter <- newRefCounter
    finPointer <- newForeignPtr context (contextFinalizer context config optionFinalizers refCounter)
    return (wrapRawContext (refCounter, finPointer), filter (/="") messages)
    where contextFinalizer context config optionFinalizers refCounter = do
            waitForZero refCounter 
            freeContext context
            freeConfig config
            sequence_ optionFinalizers

-- *Error handling
emitError :: (Context context) => context -> IO a
emitError context = do
    cs <- inContext context getError
    s <- peekCString cs
    free cs
    error s

clearError :: (Context context) => context -> IO ()
clearError context = inContext context getError >>= free

clearCache :: (Context context) => context -> IO ()
clearCache context
    = inContext context clearCaches >>= \code 
    -> if code == 0 
        then return ()
        else emitError context

-- *Syncing
sync :: (Context context) => context -> IO ()
sync context 
    = inContext context syncContext >>= \code 
    -> if code == 0 
        then return ()
        else emitError context

-- *Using the Context

-- | Do something in context and get the result.
{-# INLINE inContext #-}
inContext :: (Context c) => c -> (Ptr (RawContext c) -> IO a) -> IO a
inContext c = withForeignPtr (snd $ unwrapRawContext c) 

-- | Do something in context that emits an error code and handle it.
{-# INLINABLE inContextWithError #-}
inContextWithError :: (Context context) => context -> (Ptr (RawContext context) -> IO Int) -> IO ()
inContextWithError context f = do
    code <- attempt
    case code of
        0 -> success
        1 -> generalError
        2 -> programError
        3 -> outOfMemoryError
        _ -> unknownError
    where
        attempt = inContext context f
        success = return ()
        failure = emitError context
        generalError = failure
        programError = failure
        outOfMemoryError = do
            clearError context
            performGC
            clearCache context
            code' <- attempt
            if code' == 0
                then success
                else failure
        unknownError = failure