{-# OPTIONS_HADDOCK hide #-}
-- |
-- Module      : Data.Array.Accelerate.LLVM.Link.Cache
-- Copyright   : [2017..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.Link.Cache (

  LinkCache,
  new, dlsym,

) where

import Data.Array.Accelerate.Debug
import Data.Array.Accelerate.Lifetime
import Data.Array.Accelerate.LLVM.Compile.Cache

import Control.Monad
import Control.Concurrent.MVar
import Data.Map.Strict                                              ( Map )
import Prelude                                                      hiding ( lookup )
import Text.Printf
import qualified Data.Map.Strict                                    as Map


-- Simple reference-counted linker cache for function tables 'f' implemented by
-- object code 'o'.
--
data LinkCache f o = LinkCache {-# UNPACK #-} !(MVar (Map UID (Entry f o)))
data Entry f o     = Entry {-# UNPACK #-} !Int !f !o


-- Create a new linker cache
--
new :: IO (LinkCache f o)
new :: IO (LinkCache f o)
new = MVar (Map UID (Entry f o)) -> LinkCache f o
forall f o. MVar (Map UID (Entry f o)) -> LinkCache f o
LinkCache (MVar (Map UID (Entry f o)) -> LinkCache f o)
-> IO (MVar (Map UID (Entry f o))) -> IO (LinkCache f o)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` Map UID (Entry f o) -> IO (MVar (Map UID (Entry f o)))
forall a. a -> IO (MVar a)
newMVar Map UID (Entry f o)
forall k a. Map k a
Map.empty


-- Return the binding addresses for the given kernel functions (by key). If the
-- functions do not already exist in the cache, the supplied continuation will
-- be run in order to generate them. This happens as a single atomic step; thus
-- the cache is thread safe.
--
dlsym :: UID -> LinkCache f o -> IO (f,o) -> IO (Lifetime f)
dlsym :: UID -> LinkCache f o -> IO (f, o) -> IO (Lifetime f)
dlsym UID
key cache :: LinkCache f o
cache@(LinkCache MVar (Map UID (Entry f o))
var) IO (f, o)
k = do
  MVar (Map UID (Entry f o))
-> (Map UID (Entry f o) -> IO (Map UID (Entry f o), Lifetime f))
-> IO (Lifetime f)
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar (Map UID (Entry f o))
var ((Map UID (Entry f o) -> IO (Map UID (Entry f o), Lifetime f))
 -> IO (Lifetime f))
-> (Map UID (Entry f o) -> IO (Map UID (Entry f o), Lifetime f))
-> IO (Lifetime f)
forall a b. (a -> b) -> a -> b
$ \Map UID (Entry f o)
m ->
    case UID -> Map UID (Entry f o) -> Maybe (Entry f o)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup UID
key Map UID (Entry f o)
m of
      -- Run the supplied function to generate the object code and add to the cache
      Maybe (Entry f o)
Nothing -> do
        (f
f,o
o)  <- IO (f, o)
k
        Lifetime f
ticket <- UID -> f -> LinkCache f o -> IO (Lifetime f)
forall f o. UID -> f -> LinkCache f o -> IO (Lifetime f)
issue UID
key f
f LinkCache f o
cache
        (Map UID (Entry f o), Lifetime f)
-> IO (Map UID (Entry f o), Lifetime f)
forall (m :: * -> *) a. Monad m => a -> m a
return ( UID -> Entry f o -> Map UID (Entry f o) -> Map UID (Entry f o)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert UID
key (Int -> f -> o -> Entry f o
forall f o. Int -> f -> o -> Entry f o
Entry Int
1 f
f o
o) Map UID (Entry f o)
m, Lifetime f
ticket )

      -- Return the existing object code
      Just (Entry Int
c f
f o
o) -> do
        Lifetime f
ticket <- UID -> f -> LinkCache f o -> IO (Lifetime f)
forall f o. UID -> f -> LinkCache f o -> IO (Lifetime f)
issue UID
key f
f LinkCache f o
cache
        (Map UID (Entry f o), Lifetime f)
-> IO (Map UID (Entry f o), Lifetime f)
forall (m :: * -> *) a. Monad m => a -> m a
return ( UID -> Entry f o -> Map UID (Entry f o) -> Map UID (Entry f o)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert UID
key (Int -> f -> o -> Entry f o
forall f o. Int -> f -> o -> Entry f o
Entry (Int
cInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) f
f o
o) Map UID (Entry f o)
m, Lifetime f
ticket )


{--
-- Insert the given function table and object code into the cache. The returned
-- value must be kept alive for as long as you need the object code to live;
-- linker table entries are removed once all tickets referring to them are
-- GC'ed.
--
-- NOTE: It is an error if the entry already exists in the table. Thus, there is
-- a potential race condition between 'lookup' and 'insert'. On collision, it
-- would be fine to return a reference to the existing implementation instead
-- and discard the input values, but 'dlsym' solves this anyway.
--
insert :: Int -> f -> o -> LinkCache f o -> IO (Lifetime f)
insert key functionTable objectCode cache@(LinkCache var) = do
  ticket <- issue key functionTable cache
  modifyMVar_ var $ \m ->
    let collision = $internalError "insert" "duplicate entry"
    in  return $! Map.insertWith collision key (Entry 1 functionTable objectCode) m
  --
  return ticket


-- Check the linker cache for the given functions; if found return the
-- corresponding function table.
--
lookup :: Int -> LinkCache f o -> IO (Maybe (Lifetime f))
lookup key cache@(LinkCache var) = do
  modifyMVar var $ \m ->
    case Map.lookup key m of
      Nothing             -> return (m, Nothing)
      Just (Entry c f o)  -> do
        ticket <- issue key f cache
        return ( Map.insert key (Entry (c+1) f o) m, Just ticket )
--}


-- Issue a new ticket for the given table key/function table. When the returned
-- lifetime is GC'ed it decreasing the reference count of the corresponding
-- entry, and removes it from the table entirely once the count drops to zero.
--
issue :: UID -> f -> LinkCache f o -> IO (Lifetime f)
issue :: UID -> f -> LinkCache f o -> IO (Lifetime f)
issue UID
key f
fun (LinkCache MVar (Map UID (Entry f o))
var) = do
  Lifetime f
ticket <- f -> IO (Lifetime f)
forall a. a -> IO (Lifetime a)
newLifetime f
fun
  Lifetime f -> IO () -> IO ()
forall a. Lifetime a -> IO () -> IO ()
addFinalizer Lifetime f
ticket (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
    let refcount :: Entry f o -> Maybe (Entry f o)
refcount (Entry Int
c f
f o
o)
          | Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1    = Flag -> String -> Maybe (Entry f o) -> Maybe (Entry f o)
forall a. Flag -> String -> a -> a
trace Flag
dump_ld (String -> String -> String
forall r. PrintfType r => String -> r
printf String
"ld: remove object code %s" (UID -> String
forall a. Show a => a -> String
show UID
key)) Maybe (Entry f o)
forall a. Maybe a
Nothing
          | Bool
otherwise = Entry f o -> Maybe (Entry f o)
forall a. a -> Maybe a
Just (Int -> f -> o -> Entry f o
forall f o. Int -> f -> o -> Entry f o
Entry (Int
cInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) f
f o
o)
    in
    MVar (Map UID (Entry f o))
-> (Map UID (Entry f o) -> IO (Map UID (Entry f o))) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar (Map UID (Entry f o))
var ((Map UID (Entry f o) -> IO (Map UID (Entry f o))) -> IO ())
-> (Map UID (Entry f o) -> IO (Map UID (Entry f o))) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Map UID (Entry f o)
m -> Map UID (Entry f o) -> IO (Map UID (Entry f o))
forall (m :: * -> *) a. Monad m => a -> m a
return (Map UID (Entry f o) -> IO (Map UID (Entry f o)))
-> Map UID (Entry f o) -> IO (Map UID (Entry f o))
forall a b. (a -> b) -> a -> b
$! (Entry f o -> Maybe (Entry f o))
-> UID -> Map UID (Entry f o) -> Map UID (Entry f o)
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update Entry f o -> Maybe (Entry f o)
forall f o. Entry f o -> Maybe (Entry f o)
refcount UID
key Map UID (Entry f o)
m
  --
  Lifetime f -> IO (Lifetime f)
forall (m :: * -> *) a. Monad m => a -> m a
return Lifetime f
ticket