{-# LANGUAGE AllowAmbiguousTypes        #-}
{-# LANGUAGE BangPatterns               #-}
{-# LANGUAGE CPP                        #-}
{-# LANGUAGE ConstraintKinds            #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MagicHash                  #-}
{-# LANGUAGE PatternGuards              #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TypeApplications           #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE UnboxedTuples              #-}
{-# LANGUAGE ViewPatterns               #-}
{-# OPTIONS_HADDOCK hide #-}
-- |
-- Module      : Data.Array.Accelerate.Array.Remote.Table
-- Copyright   : [2008..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- Accelerate backends often need to copy arrays to a remote memory before they
-- can be used in computation. This module provides an automated method for
-- doing so. Keeping track of arrays in a `MemoryTable` ensures that any memory
-- allocated for them will be freed when GHC's garbage collector collects the
-- host array.
--
module Data.Array.Accelerate.Array.Remote.Table (

  -- Tables for host/device memory associations
  MemoryTable, new, lookup, malloc, free, freeStable, insertUnmanaged, reclaim,

  -- Internals
  StableArray, makeStableArray,
  makeWeakArrayData,

) where

import Control.Concurrent                                       ( yield )
import Control.Concurrent.MVar                                  ( MVar, newMVar, withMVar, mkWeakMVar )
import Control.Concurrent.Unique                                ( Unique )
import Control.Monad.IO.Class                                   ( MonadIO, liftIO )
import Data.Functor
import Data.Hashable                                            ( hash, Hashable )
import Data.Maybe                                               ( isJust )
import Data.Word
import Foreign.Storable                                         ( sizeOf )
import System.Mem                                               ( performGC )
import System.Mem.Weak                                          ( Weak, deRefWeak )
import Text.Printf
import Prelude                                                  hiding ( lookup, id )
import qualified Data.HashTable.IO                              as HT

import Data.Array.Accelerate.Error                              ( internalError )
import Data.Array.Accelerate.Type
import Data.Array.Accelerate.Array.Unique                       ( UniqueArray(..) )
import Data.Array.Accelerate.Array.Data
import Data.Array.Accelerate.Array.Remote.Class
import Data.Array.Accelerate.Array.Remote.Nursery               ( Nursery(..) )
import Data.Array.Accelerate.Lifetime
import qualified Data.Array.Accelerate.Array.Remote.Nursery     as N
import qualified Data.Array.Accelerate.Debug                    as D

import GHC.Stack


-- We use an MVar to the hash table, so that several threads may safely access
-- it concurrently. This includes the finalisation threads that remove entries
-- from the table.
--
-- It is important that we can garbage collect old entries from the table when
-- the key is no longer reachable in the heap. Hence the value part of each
-- table entry is a (Weak val), where the stable name 'key' is the key for the
-- memo table, and the 'val' is the value of this table entry. When the key
-- becomes unreachable, a finaliser will fire and remove this entry from the
-- hash buckets, and further attempts to dereference the weak pointer will
-- return Nothing. References from 'val' to the key are ignored (see the
-- semantics of weak pointers in the documentation).
--
type HashTable key val  = HT.CuckooHashTable key val
type MT p               = MVar ( HashTable StableArray (RemoteArray p) )
data MemoryTable p      = MemoryTable {-# UNPACK #-} !(MT p)
                                      {-# UNPACK #-} !(Weak (MT p))
                                      {-# UNPACK #-} !(Nursery p)
                                      (p Word8 -> IO ())

data RemoteArray p where
  RemoteArray :: !(p Word8)                 -- The actual remote pointer
              -> {-# UNPACK #-} !Int        -- The array size in bytes
              -> {-# UNPACK #-} !(Weak ())  -- Keep track of host array liveness
              -> RemoteArray p

-- | An untyped reference to an array, similar to a StableName.
--
newtype StableArray = StableArray Unique
  deriving (StableArray -> StableArray -> Bool
(StableArray -> StableArray -> Bool)
-> (StableArray -> StableArray -> Bool) -> Eq StableArray
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: StableArray -> StableArray -> Bool
$c/= :: StableArray -> StableArray -> Bool
== :: StableArray -> StableArray -> Bool
$c== :: StableArray -> StableArray -> Bool
Eq, Int -> StableArray -> Int
StableArray -> Int
(Int -> StableArray -> Int)
-> (StableArray -> Int) -> Hashable StableArray
forall a. (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: StableArray -> Int
$chash :: StableArray -> Int
hashWithSalt :: Int -> StableArray -> Int
$chashWithSalt :: Int -> StableArray -> Int
Hashable)

instance Show StableArray where
  show :: StableArray -> String
show (StableArray Unique
u) = Int -> String
forall a. Show a => a -> String
show (Unique -> Int
forall a. Hashable a => a -> Int
hash Unique
u)

-- | Create a new memory table from host to remote arrays.
--
-- The function supplied should be the `free` for the remote pointers being
-- stored. This function will be called by the GC, which typically runs on a
-- different thread. Unlike the `free` in `RemoteMemory`, this function cannot
-- depend on any state.
--
new :: (forall a. ptr a -> IO ()) -> IO (MemoryTable ptr)
new :: (forall a. ptr a -> IO ()) -> IO (MemoryTable ptr)
new forall a. ptr a -> IO ()
release = do
  String -> IO ()
forall (m :: * -> *). MonadIO m => String -> m ()
message String
"initialise memory table"
  HashTable RealWorld StableArray (RemoteArray ptr)
tbl  <- IO (HashTable RealWorld StableArray (RemoteArray ptr))
forall (h :: * -> * -> * -> *) k v.
HashTable h =>
IO (IOHashTable h k v)
HT.new
  MVar (HashTable RealWorld StableArray (RemoteArray ptr))
ref  <- HashTable RealWorld StableArray (RemoteArray ptr)
-> IO (MVar (HashTable RealWorld StableArray (RemoteArray ptr)))
forall a. a -> IO (MVar a)
newMVar HashTable RealWorld StableArray (RemoteArray ptr)
tbl
  Nursery ptr
nrs  <- (ptr Word8 -> IO ()) -> IO (Nursery ptr)
forall (ptr :: * -> *). (ptr Word8 -> IO ()) -> IO (Nursery ptr)
N.new ptr Word8 -> IO ()
forall a. ptr a -> IO ()
release
  Weak (MVar (HashTable RealWorld StableArray (RemoteArray ptr)))
weak <- MVar (HashTable RealWorld StableArray (RemoteArray ptr))
-> IO ()
-> IO
     (Weak (MVar (HashTable RealWorld StableArray (RemoteArray ptr))))
forall a. MVar a -> IO () -> IO (Weak (MVar a))
mkWeakMVar MVar (HashTable RealWorld StableArray (RemoteArray ptr))
ref (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
  MemoryTable ptr -> IO (MemoryTable ptr)
forall (m :: * -> *) a. Monad m => a -> m a
return (MemoryTable ptr -> IO (MemoryTable ptr))
-> MemoryTable ptr -> IO (MemoryTable ptr)
forall a b. (a -> b) -> a -> b
$! MT ptr
-> Weak (MT ptr)
-> Nursery ptr
-> (ptr Word8 -> IO ())
-> MemoryTable ptr
forall (p :: * -> *).
MT p
-> Weak (MT p) -> Nursery p -> (p Word8 -> IO ()) -> MemoryTable p
MemoryTable MVar (HashTable RealWorld StableArray (RemoteArray ptr))
MT ptr
ref Weak (MVar (HashTable RealWorld StableArray (RemoteArray ptr)))
Weak (MT ptr)
weak Nursery ptr
nrs ptr Word8 -> IO ()
forall a. ptr a -> IO ()
release


-- | Look for the remote pointer corresponding to a given host-side array.
--
lookup :: forall m a. (HasCallStack, RemoteMemory m)
       => MemoryTable (RemotePtr m)
       -> SingleType a
       -> ArrayData a
       -> IO (Maybe (RemotePtr m (ScalarArrayDataR a)))
lookup :: MemoryTable (RemotePtr m)
-> SingleType a
-> ArrayData a
-> IO (Maybe (RemotePtr m (ScalarArrayDataR a)))
lookup (MemoryTable !MT (RemotePtr m)
ref Weak (MT (RemotePtr m))
_ Nursery (RemotePtr m)
_ RemotePtr m Word8 -> IO ()
_) !SingleType a
tp !ArrayData a
arr
  | SingleArrayDict a
SingleArrayDict <- SingleType a -> SingleArrayDict a
forall a. SingleType a -> SingleArrayDict a
singleArrayDict SingleType a
tp = do
    StableArray
sa <- SingleType a -> ArrayData a -> IO StableArray
forall (m :: * -> *) a.
MonadIO m =>
SingleType a -> ArrayData a -> m StableArray
makeStableArray SingleType a
tp ArrayData a
arr
    Maybe (RemoteArray (RemotePtr m))
mw <- MVar (HashTable RealWorld StableArray (RemoteArray (RemotePtr m)))
-> (HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
    -> IO (Maybe (RemoteArray (RemotePtr m))))
-> IO (Maybe (RemoteArray (RemotePtr m)))
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar (HashTable RealWorld StableArray (RemoteArray (RemotePtr m)))
MT (RemotePtr m)
ref (IOHashTable HashTable StableArray (RemoteArray (RemotePtr m))
-> StableArray -> IO (Maybe (RemoteArray (RemotePtr m)))
forall (h :: * -> * -> * -> *) k v.
(HashTable h, Eq k, Hashable k) =>
IOHashTable h k v -> k -> IO (Maybe v)
`HT.lookup` StableArray
sa)
    case Maybe (RemoteArray (RemotePtr m))
mw of
      Maybe (RemoteArray (RemotePtr m))
Nothing                      -> String -> IO (Maybe (RemotePtr m a)) -> IO (Maybe (RemotePtr m a))
forall (m :: * -> *) a. MonadIO m => String -> m a -> m a
trace (String
"lookup/not found: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ StableArray -> String
forall a. Show a => a -> String
show StableArray
sa) (IO (Maybe (RemotePtr m a)) -> IO (Maybe (RemotePtr m a)))
-> IO (Maybe (RemotePtr m a)) -> IO (Maybe (RemotePtr m a))
forall a b. (a -> b) -> a -> b
$ Maybe (RemotePtr m a) -> IO (Maybe (RemotePtr m a))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (RemotePtr m a)
forall a. Maybe a
Nothing
      Just (RemoteArray RemotePtr m Word8
p Int
_ Weak ()
w) -> do
        Maybe ()
mv <- Weak () -> IO (Maybe ())
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak ()
w
        case Maybe ()
mv of
          Just{}                   -> String -> IO (Maybe (RemotePtr m a)) -> IO (Maybe (RemotePtr m a))
forall (m :: * -> *) a. MonadIO m => String -> m a -> m a
trace (String
"lookup/found: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ StableArray -> String
forall a. Show a => a -> String
show StableArray
sa) (IO (Maybe (RemotePtr m a)) -> IO (Maybe (RemotePtr m a)))
-> IO (Maybe (RemotePtr m a)) -> IO (Maybe (RemotePtr m a))
forall a b. (a -> b) -> a -> b
$ Maybe (RemotePtr m a) -> IO (Maybe (RemotePtr m a))
forall (m :: * -> *) a. Monad m => a -> m a
return (RemotePtr m a -> Maybe (RemotePtr m a)
forall a. a -> Maybe a
Just (RemotePtr m a -> Maybe (RemotePtr m a))
-> RemotePtr m a -> Maybe (RemotePtr m a)
forall a b. (a -> b) -> a -> b
$ RemotePtr m Word8 -> RemotePtr m a
forall (m :: * -> *) a b.
RemoteMemory m =>
RemotePtr m a -> RemotePtr m b
castRemotePtr @m RemotePtr m Word8
p)

          -- Note: [Weak pointer weirdness]
          --
          -- After the lookup is successful, there might conceivably be no further
          -- references to 'arr'. If that is so, and a garbage collection
          -- intervenes, the weak pointer might get tombstoned before 'deRefWeak'
          -- gets to it. In that case we throw an error (below). However, because
          -- we have used 'arr' in the continuation, this ensures that 'arr' is
          -- reachable in the continuation of 'deRefWeak' and thus 'deRefWeak'
          -- always succeeds. This sort of weirdness, typical of the world of weak
          -- pointers, is why we can not reuse the stable name 'sa' computed
          -- above in the error message.
          --
          Maybe ()
Nothing ->
            SingleType a -> ArrayData a -> IO StableArray
forall (m :: * -> *) a.
MonadIO m =>
SingleType a -> ArrayData a -> m StableArray
makeStableArray SingleType a
tp ArrayData a
arr IO StableArray
-> (StableArray -> IO (Maybe (RemotePtr m a)))
-> IO (Maybe (RemotePtr m a))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \StableArray
x -> String -> IO (Maybe (RemotePtr m a))
forall a. HasCallStack => String -> a
internalError (String -> IO (Maybe (RemotePtr m a)))
-> String -> IO (Maybe (RemotePtr m a))
forall a b. (a -> b) -> a -> b
$ String
"dead weak pair: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ StableArray -> String
forall a. Show a => a -> String
show StableArray
x

-- | Allocate a new device array to be associated with the given host-side array.
-- This may not always use the `malloc` provided by the `RemoteMemory` instance.
-- In order to reduce the number of raw allocations, previously allocated remote
-- arrays will be re-used. In the event that the remote memory is exhausted,
-- 'Nothing' is returned.
--
malloc :: forall a m. (HasCallStack, RemoteMemory m, MonadIO m)
       => MemoryTable (RemotePtr m)
       -> SingleType a
       -> ArrayData a
       -> Int
       -> m (Maybe (RemotePtr m (ScalarArrayDataR a)))
malloc :: MemoryTable (RemotePtr m)
-> SingleType a
-> ArrayData a
-> Int
-> m (Maybe (RemotePtr m (ScalarArrayDataR a)))
malloc mt :: MemoryTable (RemotePtr m)
mt@(MemoryTable MT (RemotePtr m)
_ Weak (MT (RemotePtr m))
_ !Nursery (RemotePtr m)
nursery RemotePtr m Word8 -> IO ()
_) !SingleType a
tp !ArrayData a
ad !Int
n
  | SingleArrayDict a
SingleArrayDict <- SingleType a -> SingleArrayDict a
forall a. SingleType a -> SingleArrayDict a
singleArrayDict SingleType a
tp
  , SingleDict a
SingleDict      <- SingleType a -> SingleDict a
forall a. SingleType a -> SingleDict a
singleDict SingleType a
tp
  = do
    -- Note: [Allocation sizes]
    --
    -- Instead of allocating the exact number of elements requested, we round up to
    -- a fixed chunk size as specified by RemoteMemory.remoteAllocationSize. This
    -- means there is a greater chance the nursery will get a hit, and moreover
    -- that we can search the nursery for an exact size.
    --
    Int
chunk <- m Int
forall (m :: * -> *). RemoteMemory m => m Int
remoteAllocationSize
    let -- next highest multiple of f from x
        multiple :: a -> a -> a
multiple a
x a
f      = (a
x a -> a -> a
forall a. Num a => a -> a -> a
+ (a
fa -> a -> a
forall a. Num a => a -> a -> a
-a
1)) a -> a -> a
forall a. Integral a => a -> a -> a
`quot` a
f
        bytes :: Int
bytes             = Int
chunk Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int -> Int -> Int
forall a. Integral a => a -> a -> a
multiple (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* a -> Int
forall a. Storable a => a -> Int
sizeOf (ScalarArrayDataR a
forall a. HasCallStack => a
undefined::(ScalarArrayDataR a))) Int
chunk
    --
    String -> m ()
forall (m :: * -> *). MonadIO m => String -> m ()
message (String -> m ()) -> String -> m ()
forall a b. (a -> b) -> a -> b
$ String -> Int -> Int -> Int -> String -> Int -> String
forall r. PrintfType r => String -> r
printf String
"malloc %d bytes (%d x %d bytes, type=%s, pagesize=%d)" Int
bytes Int
n (a -> Int
forall a. Storable a => a -> Int
sizeOf (ScalarArrayDataR a
forall a. HasCallStack => a
undefined:: (ScalarArrayDataR a))) (SingleType a -> String
forall a. Show a => a -> String
show SingleType a
tp) Int
chunk
    --
    Maybe (RemotePtr m a)
mp <-
      (RemotePtr m Word8 -> RemotePtr m a)
-> Maybe (RemotePtr m Word8) -> Maybe (RemotePtr m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. RemoteMemory m => RemotePtr m a -> RemotePtr m b
forall (m :: * -> *) a b.
RemoteMemory m =>
RemotePtr m a -> RemotePtr m b
castRemotePtr @m)
      (Maybe (RemotePtr m Word8) -> Maybe (RemotePtr m a))
-> m (Maybe (RemotePtr m Word8)) -> m (Maybe (RemotePtr m a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String
-> m (Maybe (RemotePtr m Word8)) -> m (Maybe (RemotePtr m Word8))
forall x. String -> m (Maybe x) -> m (Maybe x)
attempt String
"malloc/nursery" (IO (Maybe (RemotePtr m Word8)) -> m (Maybe (RemotePtr m Word8))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe (RemotePtr m Word8)) -> m (Maybe (RemotePtr m Word8)))
-> IO (Maybe (RemotePtr m Word8)) -> m (Maybe (RemotePtr m Word8))
forall a b. (a -> b) -> a -> b
$ Int -> Nursery (RemotePtr m) -> IO (Maybe (RemotePtr m Word8))
forall (ptr :: * -> *).
HasCallStack =>
Int -> Nursery ptr -> IO (Maybe (ptr Word8))
N.lookup Int
bytes Nursery (RemotePtr m)
nursery)
          m (Maybe (RemotePtr m Word8))
-> m (Maybe (RemotePtr m Word8)) -> m (Maybe (RemotePtr m Word8))
forall x. m (Maybe x) -> m (Maybe x) -> m (Maybe x)
`orElse`
          String
-> m (Maybe (RemotePtr m Word8)) -> m (Maybe (RemotePtr m Word8))
forall x. String -> m (Maybe x) -> m (Maybe x)
attempt String
"malloc/new" (Int -> m (Maybe (RemotePtr m Word8))
forall (m :: * -> *).
RemoteMemory m =>
Int -> m (Maybe (RemotePtr m Word8))
mallocRemote Int
bytes)
          m (Maybe (RemotePtr m Word8))
-> m (Maybe (RemotePtr m Word8)) -> m (Maybe (RemotePtr m Word8))
forall x. m (Maybe x) -> m (Maybe x) -> m (Maybe x)
`orElse` do String -> m ()
forall (m :: * -> *). MonadIO m => String -> m ()
message String
"malloc/remote-malloc-failed (cleaning)"
                      MemoryTable (RemotePtr m) -> m ()
forall (m :: * -> *).
(RemoteMemory m, MonadIO m) =>
MemoryTable (RemotePtr m) -> m ()
clean MemoryTable (RemotePtr m)
mt
                      IO (Maybe (RemotePtr m Word8)) -> m (Maybe (RemotePtr m Word8))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe (RemotePtr m Word8)) -> m (Maybe (RemotePtr m Word8)))
-> IO (Maybe (RemotePtr m Word8)) -> m (Maybe (RemotePtr m Word8))
forall a b. (a -> b) -> a -> b
$ Int -> Nursery (RemotePtr m) -> IO (Maybe (RemotePtr m Word8))
forall (ptr :: * -> *).
HasCallStack =>
Int -> Nursery ptr -> IO (Maybe (ptr Word8))
N.lookup Int
bytes Nursery (RemotePtr m)
nursery
          m (Maybe (RemotePtr m Word8))
-> m (Maybe (RemotePtr m Word8)) -> m (Maybe (RemotePtr m Word8))
forall x. m (Maybe x) -> m (Maybe x) -> m (Maybe x)
`orElse` do String -> m ()
forall (m :: * -> *). MonadIO m => String -> m ()
message String
"malloc/remote-malloc-failed (purging)"
                      MemoryTable (RemotePtr m) -> m ()
forall (m :: * -> *).
(RemoteMemory m, MonadIO m) =>
MemoryTable (RemotePtr m) -> m ()
purge MemoryTable (RemotePtr m)
mt
                      Int -> m (Maybe (RemotePtr m Word8))
forall (m :: * -> *).
RemoteMemory m =>
Int -> m (Maybe (RemotePtr m Word8))
mallocRemote Int
bytes
          m (Maybe (RemotePtr m Word8))
-> m (Maybe (RemotePtr m Word8)) -> m (Maybe (RemotePtr m Word8))
forall x. m (Maybe x) -> m (Maybe x) -> m (Maybe x)
`orElse` do String -> m ()
forall (m :: * -> *). MonadIO m => String -> m ()
message String
"malloc/remote-malloc-failed (non-recoverable)"
                      Maybe (RemotePtr m Word8) -> m (Maybe (RemotePtr m Word8))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (RemotePtr m Word8)
forall a. Maybe a
Nothing
    case Maybe (RemotePtr m a)
mp of
      Maybe (RemotePtr m a)
Nothing -> Maybe (RemotePtr m a) -> m (Maybe (RemotePtr m a))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (RemotePtr m a)
forall a. Maybe a
Nothing
      Just RemotePtr m a
p' -> do
        MemoryTable (RemotePtr m)
-> SingleType a
-> ArrayData a
-> RemotePtr m (ScalarArrayDataR a)
-> Int
-> m ()
forall (m :: * -> *) a.
(RemoteMemory m, MonadIO m) =>
MemoryTable (RemotePtr m)
-> SingleType a
-> ArrayData a
-> RemotePtr m (ScalarArrayDataR a)
-> Int
-> m ()
insert MemoryTable (RemotePtr m)
mt SingleType a
tp ArrayData a
ad RemotePtr m a
RemotePtr m (ScalarArrayDataR a)
p' Int
bytes
        Maybe (RemotePtr m a) -> m (Maybe (RemotePtr m a))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (RemotePtr m a)
mp
  where
    {-# INLINE orElse #-}
    orElse :: m (Maybe x) -> m (Maybe x) -> m (Maybe x)
    orElse :: m (Maybe x) -> m (Maybe x) -> m (Maybe x)
orElse m (Maybe x)
this m (Maybe x)
next = do
      Maybe x
result <- m (Maybe x)
this
      case Maybe x
result of
        Just{}  -> Maybe x -> m (Maybe x)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe x
result
        Maybe x
Nothing -> m (Maybe x)
next

    {-# INLINE attempt #-}
    attempt :: String -> m (Maybe x) -> m (Maybe x)
    attempt :: String -> m (Maybe x) -> m (Maybe x)
attempt String
msg m (Maybe x)
this = do
      Maybe x
result <- m (Maybe x)
this
      case Maybe x
result of
        Just{}  -> String -> m (Maybe x) -> m (Maybe x)
forall (m :: * -> *) a. MonadIO m => String -> m a -> m a
trace String
msg (Maybe x -> m (Maybe x)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe x
result)
        Maybe x
Nothing -> Maybe x -> m (Maybe x)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe x
forall a. Maybe a
Nothing



-- | Deallocate the device array associated with the given host-side array.
-- Typically this should only be called in very specific circumstances.
--
free :: forall m a. (RemoteMemory m)
     => MemoryTable (RemotePtr m)
     -> SingleType a
     -> ArrayData a
     -> IO ()
free :: MemoryTable (RemotePtr m) -> SingleType a -> ArrayData a -> IO ()
free MemoryTable (RemotePtr m)
mt SingleType a
tp !ArrayData a
arr = do
  StableArray
sa <- SingleType a -> ArrayData a -> IO StableArray
forall (m :: * -> *) a.
MonadIO m =>
SingleType a -> ArrayData a -> m StableArray
makeStableArray SingleType a
tp ArrayData a
arr
  MemoryTable (RemotePtr m) -> StableArray -> IO ()
forall (m :: * -> *).
RemoteMemory m =>
MemoryTable (RemotePtr m) -> StableArray -> IO ()
freeStable @m MemoryTable (RemotePtr m)
mt StableArray
sa


-- | Deallocate the device array associated with the given StableArray. This
-- is useful for other memory managers built on top of the memory table.
--
freeStable
    :: forall m. RemoteMemory m
    => MemoryTable (RemotePtr m)
    -> StableArray
    -> IO ()
freeStable :: MemoryTable (RemotePtr m) -> StableArray -> IO ()
freeStable (MemoryTable !MT (RemotePtr m)
ref Weak (MT (RemotePtr m))
_ !Nursery (RemotePtr m)
nrs RemotePtr m Word8 -> IO ()
_) !StableArray
sa =
  MVar (HashTable RealWorld StableArray (RemoteArray (RemotePtr m)))
-> (HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
    -> IO ())
-> IO ()
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar (HashTable RealWorld StableArray (RemoteArray (RemotePtr m)))
MT (RemotePtr m)
ref      ((HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
  -> IO ())
 -> IO ())
-> (HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
    -> IO ())
-> IO ()
forall a b. (a -> b) -> a -> b
$ \HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
mt ->
  IOHashTable HashTable StableArray (RemoteArray (RemotePtr m))
-> StableArray
-> (Maybe (RemoteArray (RemotePtr m))
    -> IO (Maybe (RemoteArray (RemotePtr m)), ()))
-> IO ()
forall (h :: * -> * -> * -> *) k v a.
(HashTable h, Eq k, Hashable k) =>
IOHashTable h k v -> k -> (Maybe v -> IO (Maybe v, a)) -> IO a
HT.mutateIO HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
IOHashTable HashTable StableArray (RemoteArray (RemotePtr m))
mt StableArray
sa ((Maybe (RemoteArray (RemotePtr m))
  -> IO (Maybe (RemoteArray (RemotePtr m)), ()))
 -> IO ())
-> (Maybe (RemoteArray (RemotePtr m))
    -> IO (Maybe (RemoteArray (RemotePtr m)), ()))
-> IO ()
forall a b. (a -> b) -> a -> b
$ \Maybe (RemoteArray (RemotePtr m))
mw -> do
    case Maybe (RemoteArray (RemotePtr m))
mw of
      Maybe (RemoteArray (RemotePtr m))
Nothing ->
        String -> IO ()
forall (m :: * -> *). MonadIO m => String -> m ()
message (String
"free/already-removed: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ StableArray -> String
forall a. Show a => a -> String
show StableArray
sa)

      Just (RemoteArray !RemotePtr m Word8
p !Int
bytes Weak ()
_) -> do
        String -> IO ()
forall (m :: * -> *). MonadIO m => String -> m ()
message (String
"free/nursery: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ StableArray -> String
forall a. Show a => a -> String
show StableArray
sa String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" of " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall n. Integral n => n -> String
showBytes Int
bytes)
        Int -> RemotePtr m Word8 -> Nursery (RemotePtr m) -> IO ()
forall (ptr :: * -> *). Int -> ptr Word8 -> Nursery ptr -> IO ()
N.insert Int
bytes (RemotePtr m Word8 -> RemotePtr m Word8
forall (m :: * -> *) a b.
RemoteMemory m =>
RemotePtr m a -> RemotePtr m b
castRemotePtr @m RemotePtr m Word8
p) Nursery (RemotePtr m)
nrs
        Int64 -> IO ()
D.decreaseCurrentBytesRemote (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
bytes)

    (Maybe (RemoteArray (RemotePtr m)), ())
-> IO (Maybe (RemoteArray (RemotePtr m)), ())
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (RemoteArray (RemotePtr m))
forall a. Maybe a
Nothing, ())


-- | Record an association between a host-side array and a new device memory
-- area. The device memory will be freed when the host array is garbage
-- collected.
--
insert
    :: forall m a. (RemoteMemory m, MonadIO m)
    => MemoryTable (RemotePtr m)
    -> SingleType a
    -> ArrayData a
    -> RemotePtr m (ScalarArrayDataR a)
    -> Int
    -> m ()
insert :: MemoryTable (RemotePtr m)
-> SingleType a
-> ArrayData a
-> RemotePtr m (ScalarArrayDataR a)
-> Int
-> m ()
insert mt :: MemoryTable (RemotePtr m)
mt@(MemoryTable !MT (RemotePtr m)
ref Weak (MT (RemotePtr m))
_ Nursery (RemotePtr m)
_ RemotePtr m Word8 -> IO ()
_) !SingleType a
tp !ArrayData a
arr !RemotePtr m (ScalarArrayDataR a)
ptr !Int
bytes | SingleArrayDict a
SingleArrayDict <- SingleType a -> SingleArrayDict a
forall a. SingleType a -> SingleArrayDict a
singleArrayDict SingleType a
tp = do
  StableArray
key  <- SingleType a -> ArrayData a -> m StableArray
forall (m :: * -> *) a.
MonadIO m =>
SingleType a -> ArrayData a -> m StableArray
makeStableArray SingleType a
tp ArrayData a
arr
  Weak ()
weak <- IO (Weak ()) -> m (Weak ())
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Weak ()) -> m (Weak ())) -> IO (Weak ()) -> m (Weak ())
forall a b. (a -> b) -> a -> b
$ SingleType a -> ArrayData a -> () -> Maybe (IO ()) -> IO (Weak ())
forall e c.
SingleType e -> ArrayData e -> c -> Maybe (IO ()) -> IO (Weak c)
makeWeakArrayData SingleType a
tp ArrayData a
arr () (IO () -> Maybe (IO ())
forall a. a -> Maybe a
Just (IO () -> Maybe (IO ())) -> IO () -> Maybe (IO ())
forall a b. (a -> b) -> a -> b
$ MemoryTable (RemotePtr m) -> StableArray -> IO ()
forall (m :: * -> *).
RemoteMemory m =>
MemoryTable (RemotePtr m) -> StableArray -> IO ()
freeStable @m MemoryTable (RemotePtr m)
mt StableArray
key)
  String -> m ()
forall (m :: * -> *). MonadIO m => String -> m ()
message (String -> m ()) -> String -> m ()
forall a b. (a -> b) -> a -> b
$ String
"insert: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ StableArray -> String
forall a. Show a => a -> String
show StableArray
key
  IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO  (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Int64 -> IO ()
D.increaseCurrentBytesRemote (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
bytes)
  IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO  (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ MVar (HashTable RealWorld StableArray (RemoteArray (RemotePtr m)))
-> (HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
    -> IO ())
-> IO ()
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar (HashTable RealWorld StableArray (RemoteArray (RemotePtr m)))
MT (RemotePtr m)
ref ((HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
  -> IO ())
 -> IO ())
-> (HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
    -> IO ())
-> IO ()
forall a b. (a -> b) -> a -> b
$ \HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
tbl -> IOHashTable HashTable StableArray (RemoteArray (RemotePtr m))
-> StableArray -> RemoteArray (RemotePtr m) -> IO ()
forall (h :: * -> * -> * -> *) k v.
(HashTable h, Eq k, Hashable k) =>
IOHashTable h k v -> k -> v -> IO ()
HT.insert HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
IOHashTable HashTable StableArray (RemoteArray (RemotePtr m))
tbl StableArray
key (RemotePtr m Word8 -> Int -> Weak () -> RemoteArray (RemotePtr m)
forall (p :: * -> *). p Word8 -> Int -> Weak () -> RemoteArray p
RemoteArray (RemotePtr m a -> RemotePtr m Word8
forall (m :: * -> *) a b.
RemoteMemory m =>
RemotePtr m a -> RemotePtr m b
castRemotePtr @m RemotePtr m a
RemotePtr m (ScalarArrayDataR a)
ptr) Int
bytes Weak ()
weak)


-- | Record an association between a host-side array and a remote memory area
-- that was not allocated by accelerate. The remote memory will NOT be re-used
-- once the host-side array is garbage collected.
--
-- This typically only has use for backends that provide an FFI.
--
insertUnmanaged
    :: forall m a. (MonadIO m, RemoteMemory m)
    => MemoryTable (RemotePtr m)
    -> SingleType a
    -> ArrayData a
    -> RemotePtr m (ScalarArrayDataR a)
    -> m ()
insertUnmanaged :: MemoryTable (RemotePtr m)
-> SingleType a
-> ArrayData a
-> RemotePtr m (ScalarArrayDataR a)
-> m ()
insertUnmanaged (MemoryTable !MT (RemotePtr m)
ref !Weak (MT (RemotePtr m))
weak_ref Nursery (RemotePtr m)
_ RemotePtr m Word8 -> IO ()
_) SingleType a
tp !ArrayData a
arr !RemotePtr m (ScalarArrayDataR a)
ptr | SingleArrayDict a
SingleArrayDict  <- SingleType a -> SingleArrayDict a
forall a. SingleType a -> SingleArrayDict a
singleArrayDict SingleType a
tp = do
  StableArray
key  <- SingleType a -> ArrayData a -> m StableArray
forall (m :: * -> *) a.
MonadIO m =>
SingleType a -> ArrayData a -> m StableArray
makeStableArray SingleType a
tp ArrayData a
arr
  Weak ()
weak <- IO (Weak ()) -> m (Weak ())
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Weak ()) -> m (Weak ())) -> IO (Weak ()) -> m (Weak ())
forall a b. (a -> b) -> a -> b
$ SingleType a -> ArrayData a -> () -> Maybe (IO ()) -> IO (Weak ())
forall e c.
SingleType e -> ArrayData e -> c -> Maybe (IO ()) -> IO (Weak c)
makeWeakArrayData SingleType a
tp ArrayData a
arr () (IO () -> Maybe (IO ())
forall a. a -> Maybe a
Just (IO () -> Maybe (IO ())) -> IO () -> Maybe (IO ())
forall a b. (a -> b) -> a -> b
$ Weak (MT (RemotePtr m)) -> StableArray -> IO ()
forall (p :: * -> *). Weak (MT p) -> StableArray -> IO ()
remoteFinalizer Weak (MT (RemotePtr m))
weak_ref StableArray
key)
  String -> m ()
forall (m :: * -> *). MonadIO m => String -> m ()
message (String -> m ()) -> String -> m ()
forall a b. (a -> b) -> a -> b
$ String
"insertUnmanaged: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ StableArray -> String
forall a. Show a => a -> String
show StableArray
key
  IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO  (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ MVar (HashTable RealWorld StableArray (RemoteArray (RemotePtr m)))
-> (HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
    -> IO ())
-> IO ()
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar (HashTable RealWorld StableArray (RemoteArray (RemotePtr m)))
MT (RemotePtr m)
ref ((HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
  -> IO ())
 -> IO ())
-> (HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
    -> IO ())
-> IO ()
forall a b. (a -> b) -> a -> b
$ \HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
tbl -> IOHashTable HashTable StableArray (RemoteArray (RemotePtr m))
-> StableArray -> RemoteArray (RemotePtr m) -> IO ()
forall (h :: * -> * -> * -> *) k v.
(HashTable h, Eq k, Hashable k) =>
IOHashTable h k v -> k -> v -> IO ()
HT.insert HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
IOHashTable HashTable StableArray (RemoteArray (RemotePtr m))
tbl StableArray
key (RemotePtr m Word8 -> Int -> Weak () -> RemoteArray (RemotePtr m)
forall (p :: * -> *). p Word8 -> Int -> Weak () -> RemoteArray p
RemoteArray (RemotePtr m a -> RemotePtr m Word8
forall (m :: * -> *) a b.
RemoteMemory m =>
RemotePtr m a -> RemotePtr m b
castRemotePtr @m RemotePtr m a
RemotePtr m (ScalarArrayDataR a)
ptr) Int
0 Weak ()
weak)


-- Removing entries
-- ----------------

-- | Initiate garbage collection and mark any arrays that no longer have
-- host-side equivalents as reusable.
--
clean :: forall m. (RemoteMemory m, MonadIO m) => MemoryTable (RemotePtr m) -> m ()
clean :: MemoryTable (RemotePtr m) -> m ()
clean mt :: MemoryTable (RemotePtr m)
mt@(MemoryTable MT (RemotePtr m)
_ Weak (MT (RemotePtr m))
weak_ref Nursery (RemotePtr m)
nrs RemotePtr m Word8 -> IO ()
_) = String -> Nursery (RemotePtr m) -> m () -> m ()
forall (m :: * -> *) (p :: * -> *) a.
(RemoteMemory m, MonadIO m) =>
String -> Nursery p -> m a -> m a
management String
"clean" Nursery (RemotePtr m)
nrs (m () -> m ()) -> (IO () -> m ()) -> IO () -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
  -- Unfortunately there is no real way to force a GC then wait for it to
  -- finish. Calling performGC then yielding works moderately well in
  -- single-threaded cases, but tends to fall down otherwise. Either way, given
  -- that finalizers are often significantly delayed, it is worth our while
  -- traversing the table and explicitly freeing any dead entires.
  --
  IO ()
D.didRemoteGC
  IO ()
performGC
  IO ()
yield
  Maybe
  (MVar
     (HashTable RealWorld StableArray (RemoteArray (RemotePtr m))))
mr <- Weak
  (MVar
     (HashTable RealWorld StableArray (RemoteArray (RemotePtr m))))
-> IO
     (Maybe
        (MVar
           (HashTable RealWorld StableArray (RemoteArray (RemotePtr m)))))
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak
  (MVar
     (HashTable RealWorld StableArray (RemoteArray (RemotePtr m))))
Weak (MT (RemotePtr m))
weak_ref
  case Maybe
  (MVar
     (HashTable RealWorld StableArray (RemoteArray (RemotePtr m))))
mr of
    Maybe
  (MVar
     (HashTable RealWorld StableArray (RemoteArray (RemotePtr m))))
Nothing  -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    Just MVar (HashTable RealWorld StableArray (RemoteArray (RemotePtr m)))
ref -> do
      [StableArray]
rs <- MVar (HashTable RealWorld StableArray (RemoteArray (RemotePtr m)))
-> (HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
    -> IO [StableArray])
-> IO [StableArray]
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar (HashTable RealWorld StableArray (RemoteArray (RemotePtr m)))
ref ((HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
  -> IO [StableArray])
 -> IO [StableArray])
-> (HashTable RealWorld StableArray (RemoteArray (RemotePtr m))
    -> IO [StableArray])
-> IO [StableArray]
forall a b. (a -> b) -> a -> b
$ ([StableArray]
 -> (StableArray, RemoteArray (RemotePtr m)) -> IO [StableArray])
-> [StableArray]
-> IOHashTable HashTable StableArray (RemoteArray (RemotePtr m))
-> IO [StableArray]
forall (h :: * -> * -> * -> *) a k v.
HashTable h =>
(a -> (k, v) -> IO a) -> a -> IOHashTable h k v -> IO a
HT.foldM [StableArray]
-> (StableArray, RemoteArray (RemotePtr m)) -> IO [StableArray]
forall a (p :: * -> *). [a] -> (a, RemoteArray p) -> IO [a]
removable []  -- collect arrays that can be removed
      (StableArray -> IO ()) -> [StableArray] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (MemoryTable (RemotePtr m) -> StableArray -> IO ()
forall (m :: * -> *).
RemoteMemory m =>
MemoryTable (RemotePtr m) -> StableArray -> IO ()
freeStable @m MemoryTable (RemotePtr m)
mt) [StableArray]
rs -- remove them all
  where
    removable :: [a] -> (a, RemoteArray p) -> IO [a]
removable [a]
rs (a
sa, RemoteArray p Word8
_ Int
_ Weak ()
w) = do
      Bool
alive <- Maybe () -> Bool
forall a. Maybe a -> Bool
isJust (Maybe () -> Bool) -> IO (Maybe ()) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Weak () -> IO (Maybe ())
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak ()
w
      if Bool
alive
        then [a] -> IO [a]
forall (m :: * -> *) a. Monad m => a -> m a
return [a]
rs
        else [a] -> IO [a]
forall (m :: * -> *) a. Monad m => a -> m a
return (a
saa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
rs)


-- | Call `free` on all arrays that are not currently associated with host-side
-- arrays.
--
purge :: (RemoteMemory m, MonadIO m) => MemoryTable (RemotePtr m) -> m ()
purge :: MemoryTable (RemotePtr m) -> m ()
purge (MemoryTable MT (RemotePtr m)
_ Weak (MT (RemotePtr m))
_ nursery :: Nursery (RemotePtr m)
nursery@(Nursery NRS (RemotePtr m)
nrs Weak (NRS (RemotePtr m))
_) RemotePtr m Word8 -> IO ()
release)
  = String -> Nursery (RemotePtr m) -> m () -> m ()
forall (m :: * -> *) (p :: * -> *) a.
(RemoteMemory m, MonadIO m) =>
String -> Nursery p -> m a -> m a
management String
"purge" Nursery (RemotePtr m)
nursery
  (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO ((RemotePtr m Word8 -> IO ()) -> NRS (RemotePtr m) -> IO ()
forall (ptr :: * -> *). (ptr Word8 -> IO ()) -> NRS ptr -> IO ()
N.cleanup RemotePtr m Word8 -> IO ()
release NRS (RemotePtr m)
nrs)


-- | Initiate garbage collection and `free` any remote arrays that no longer
-- have matching host-side equivalents.
--
reclaim :: forall m. (RemoteMemory m, MonadIO m) => MemoryTable (RemotePtr m) -> m ()
reclaim :: MemoryTable (RemotePtr m) -> m ()
reclaim MemoryTable (RemotePtr m)
mt = MemoryTable (RemotePtr m) -> m ()
forall (m :: * -> *).
(RemoteMemory m, MonadIO m) =>
MemoryTable (RemotePtr m) -> m ()
clean MemoryTable (RemotePtr m)
mt m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MemoryTable (RemotePtr m) -> m ()
forall (m :: * -> *).
(RemoteMemory m, MonadIO m) =>
MemoryTable (RemotePtr m) -> m ()
purge MemoryTable (RemotePtr m)
mt

remoteFinalizer :: Weak (MT p) -> StableArray -> IO ()
remoteFinalizer :: Weak (MT p) -> StableArray -> IO ()
remoteFinalizer !Weak (MT p)
weak_ref !StableArray
key = do
  Maybe (MVar (HashTable RealWorld StableArray (RemoteArray p)))
mr <- Weak (MVar (HashTable RealWorld StableArray (RemoteArray p)))
-> IO
     (Maybe (MVar (HashTable RealWorld StableArray (RemoteArray p))))
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak (MVar (HashTable RealWorld StableArray (RemoteArray p)))
Weak (MT p)
weak_ref
  case Maybe (MVar (HashTable RealWorld StableArray (RemoteArray p)))
mr of
    Maybe (MVar (HashTable RealWorld StableArray (RemoteArray p)))
Nothing  -> String -> IO ()
forall (m :: * -> *). MonadIO m => String -> m ()
message (String
"finalise/dead table: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ StableArray -> String
forall a. Show a => a -> String
show StableArray
key)
    Just MVar (HashTable RealWorld StableArray (RemoteArray p))
ref -> String -> IO () -> IO ()
forall (m :: * -> *) a. MonadIO m => String -> m a -> m a
trace   (String
"finalise: "            String -> ShowS
forall a. [a] -> [a] -> [a]
++ StableArray -> String
forall a. Show a => a -> String
show StableArray
key) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ MVar (HashTable RealWorld StableArray (RemoteArray p))
-> (HashTable RealWorld StableArray (RemoteArray p) -> IO ())
-> IO ()
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar (HashTable RealWorld StableArray (RemoteArray p))
ref (IOHashTable HashTable StableArray (RemoteArray p)
-> StableArray -> IO ()
forall (h :: * -> * -> * -> *) k v.
(HashTable h, Eq k, Hashable k) =>
IOHashTable h k v -> k -> IO ()
`HT.delete` StableArray
key)


-- Miscellaneous
-- -------------

-- | Make a new 'StableArray'.
--
{-# INLINE makeStableArray #-}
makeStableArray
    :: MonadIO m
    => SingleType a
    -> ArrayData a
    -> m StableArray
makeStableArray :: SingleType a -> ArrayData a -> m StableArray
makeStableArray !SingleType a
tp !ArrayData a
ad
  | SingleArrayDict a
SingleArrayDict <- SingleType a -> SingleArrayDict a
forall a. SingleType a -> SingleArrayDict a
singleArrayDict SingleType a
tp
  = StableArray -> m StableArray
forall (m :: * -> *) a. Monad m => a -> m a
return (StableArray -> m StableArray) -> StableArray -> m StableArray
forall a b. (a -> b) -> a -> b
$! Unique -> StableArray
StableArray (UniqueArray a -> Unique
forall e. UniqueArray e -> Unique
uniqueArrayId UniqueArray a
ArrayData a
ad)


-- Weak arrays
-- -----------

-- | Make a weak pointer using an array as a key. Unlike the standard `mkWeak`,
-- this guarantees finalisers won't fire early.
--
makeWeakArrayData
    :: forall e c.
       SingleType e
    -> ArrayData e
    -> c
    -> Maybe (IO ())
    -> IO (Weak c)
makeWeakArrayData :: SingleType e -> ArrayData e -> c -> Maybe (IO ()) -> IO (Weak c)
makeWeakArrayData !SingleType e
tp !ArrayData e
ad !c
c !Maybe (IO ())
mf | SingleArrayDict e
SingleArrayDict <- SingleType e -> SingleArrayDict e
forall a. SingleType a -> SingleArrayDict a
singleArrayDict SingleType e
tp = do
  let !uad :: Lifetime (ForeignPtr e)
uad = UniqueArray e -> Lifetime (ForeignPtr e)
forall e. UniqueArray e -> Lifetime (ForeignPtr e)
uniqueArrayData UniqueArray e
ArrayData e
ad
  case Maybe (IO ())
mf of
    Maybe (IO ())
Nothing -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    Just IO ()
f  -> Lifetime (ForeignPtr e) -> IO () -> IO ()
forall a. Lifetime a -> IO () -> IO ()
addFinalizer Lifetime (ForeignPtr e)
uad IO ()
f
  Lifetime (ForeignPtr e) -> c -> IO (Weak c)
forall k v. Lifetime k -> v -> IO (Weak v)
mkWeak Lifetime (ForeignPtr e)
uad c
c


-- Debug
-- -----

{-# INLINE showBytes #-}
showBytes :: Integral n => n -> String
showBytes :: n -> String
showBytes n
x = Maybe Int -> Double -> Double -> ShowS
forall a. RealFloat a => Maybe Int -> a -> a -> ShowS
D.showFFloatSIBase (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
0) Double
1024 (n -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral n
x :: Double) String
"B"

{-# INLINE trace #-}
trace :: MonadIO m => String -> m a -> m a
trace :: String -> m a -> m a
trace String
msg m a
next = String -> m ()
forall (m :: * -> *). MonadIO m => String -> m ()
message String
msg m () -> m a -> m a
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m a
next

{-# INLINE message #-}
message :: MonadIO m => String -> m ()
message :: String -> m ()
message String
msg = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Flag -> String -> IO ()
D.traceIO Flag
D.dump_gc (String
"gc: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
msg)

{-# INLINE management #-}
management :: (RemoteMemory m, MonadIO m) => String -> Nursery p -> m a -> m a
management :: String -> Nursery p -> m a -> m a
management String
msg Nursery p
nrs m a
next = do
  Bool
yes <- IO Bool -> m Bool
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Bool -> m Bool) -> IO Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ Flag -> IO Bool
D.getFlag Flag
D.dump_gc
  if Bool
yes
    then do
      Int64
total       <- m Int64
forall (m :: * -> *). RemoteMemory m => m Int64
totalRemoteMem
      Int64
before      <- m Int64
forall (m :: * -> *). RemoteMemory m => m Int64
availableRemoteMem
      Int64
before_nrs  <- IO Int64 -> m Int64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int64 -> m Int64) -> IO Int64 -> m Int64
forall a b. (a -> b) -> a -> b
$ Nursery p -> IO Int64
forall (ptr :: * -> *). Nursery ptr -> IO Int64
N.size Nursery p
nrs
      a
r           <- m a
next
      Int64
after       <- m Int64
forall (m :: * -> *). RemoteMemory m => m Int64
availableRemoteMem
      Int64
after_nrs   <- IO Int64 -> m Int64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int64 -> m Int64) -> IO Int64 -> m Int64
forall a b. (a -> b) -> a -> b
$ Nursery p -> IO Int64
forall (ptr :: * -> *). Nursery ptr -> IO Int64
N.size Nursery p
nrs
      String -> m ()
forall (m :: * -> *). MonadIO m => String -> m ()
message (String -> m ()) -> String -> m ()
forall a b. (a -> b) -> a -> b
$ String -> String -> String -> String -> String -> ShowS
forall r. PrintfType r => String -> r
printf String
"%s (freed: %s, stashed: %s, remaining: %s of %s)"
                  String
msg
                  (Int64 -> String
forall n. Integral n => n -> String
showBytes (Int64
before Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
after))
                  (Int64 -> String
forall n. Integral n => n -> String
showBytes (Int64
after_nrs Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
before_nrs))
                  (Int64 -> String
forall n. Integral n => n -> String
showBytes Int64
after)
                  (Int64 -> String
forall n. Integral n => n -> String
showBytes Int64
total)
      --
      a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
    else
      m a
next