Change Nursery representation

This changes from a Hashtable of FullList to an IntMap of Sequence. I should benchmark this properly, but from experiments a long time ago this was faster. The improved IntMap API also allows for fewer querys (e.g. combining lookup and update).

Data/Array/Accelerate/Array/Remote/Nursery.hs

@@ -1,10 +1,5 @@
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeFamilies #-}
-{-# LANGUAGE ViewPatterns #-}
-{-# LANGUAGE PolyKinds #-}
-{-# OPTIONS_GHC -fno-warn-orphans #-}
+{-# LANGUAGE BangPatterns    #-}
+{-# LANGUAGE TemplateHaskell #-}
 -- |
 -- Module      : Data.Array.Accelerate.Array.Remote.Nursery
 -- Copyright   : [2008..2014] Manuel M T Chakravarty, Gabriele Keller
@@ -19,103 +14,103 @@
 
 module Data.Array.Accelerate.Array.Remote.Nursery (
 
-  Nursery(..), NRS, new, malloc, stash, flush, size
+  Nursery(..), NRS, new, lookup, insert, cleanup, size
 
 ) where
 
 -- friends
-import Data.Array.Accelerate.FullList                           ( FullList(..) )
-import Data.Array.Accelerate.Array.Remote.Class
-import qualified Data.Array.Accelerate.FullList                 as FL
+import Data.Array.Accelerate.Error
 import qualified Data.Array.Accelerate.Debug                    as D
 
 -- libraries
+import Prelude                                                  hiding ( lookup )
 import Control.Concurrent.MVar
 import Data.Int
-import Data.Proxy
+import Data.IntMap                                              ( IntMap )
+import Data.Sequence                                            ( Seq )
 import Data.Word
 import System.Mem.Weak                                          ( Weak )
-import Prelude
-
-import qualified Data.HashTable.IO                              as HT
+import qualified Data.IntMap.Strict                             as IM
+import qualified Data.Sequence                                  as Seq
+import qualified Data.Traversable                               as Seq
 
 
 -- The nursery is a place to store remote memory arrays that are no longer
--- needed. If a new array is requested of a similar size, we might return an
--- array from the nursery instead of calling into the backends underlying API
--- to allocate fresh memory.
+-- needed. Often it is quicker to reuse an existing array, rather than call out
+-- to the external API to allocate fresh memory.
 --
--- Note that since there might be many arrays for the same size, each entry in
--- the map keeps a (non-empty) list of remote arrays.
+-- The nursery is wrapped in an MVar so that several threads may safely access
+-- it concurrently.
 --
-type HashTable key val = HT.BasicHashTable key val
+data Nursery ptr        = Nursery {-# UNPACK #-} !(NRS ptr)
+                                  {-# UNPACK #-} !(Weak (NRS ptr))
+type NRS ptr            = MVar (N ptr)
 
-type NRS ptr           = MVar ( HashTable Int (FullList () (ptr Word8)), Int64 )
-data Nursery ptr       = Nursery {-# UNPACK #-} !(NRS ptr)
-                                 {-# UNPACK #-} !(Weak (NRS ptr))
+data N ptr              = N !(IntMap (Seq (ptr Word8)))       -- #bytes -> ptr
+                            {-# UNPACK #-} !Int64             -- total allocated bytes
 
 
--- Generate a fresh nursery
+-- | Create a fresh nursery.
+--
+-- When the nursery is garbage collected, the provided function will be run on
+-- each value to free the retained memory.
 --
+{-# INLINEABLE new #-}
 new :: (ptr Word8 -> IO ()) -> IO (Nursery ptr)
-new free = do
-  tbl    <- HT.new
-  ref    <- newMVar (tbl, 0)
-  weak   <- mkWeakMVar ref (flush free tbl)
+new delete = do
+  message "initialise nursery"
+  ref    <- newMVar ( N IM.empty 0 )
+  weak   <- mkWeakMVar ref (cleanup delete ref)
   return $! Nursery ref weak
 
 
--- Look for a chunk of memory in the nursery of a given size (or a little bit
--- larger). If found, it is removed from the nursery and a pointer to it
--- returned.
+-- | Look for an entry with the requested size.
 --
-{-# INLINE malloc #-}
-malloc :: Int
-       -> Nursery ptr
-       -> IO (Maybe (ptr Word8))
-malloc !n (Nursery !ref _) = modifyMVar ref $ \(tbl,sz) -> do
-  mp  <- HT.lookup tbl n
-  case mp of
-    Nothing               -> return ((tbl,sz),Nothing)
-    Just (FL () ptr rest) ->
-      case rest of
-        FL.Nil          -> HT.delete tbl n              >> return ((tbl,sz - fromIntegral n), Just ptr)
-        FL.Cons () v xs -> HT.insert tbl n (FL () v xs) >> return ((tbl,sz - fromIntegral n), Just ptr)
-
-
--- Add a device pointer to the nursery.
+{-# INLINEABLE lookup #-}
+lookup :: Int -> Nursery ptr -> IO (Maybe (ptr Word8))
+lookup !key (Nursery !ref !_) =
+  modifyMVar ref $ \nrs@( N im sz ) ->
+    let
+        (mv, nrs')  = IM.updateLookupWithKey f key im         -- returns _original_ value, if located
+        f _k v      =
+          case Seq.viewl v of
+            Seq.EmptyL  -> $internalError "lookup" "expected non-empty sequence"
+            _ Seq.:< vs -> if Seq.null vs then Nothing        -- delete this entry in the map
+                                          else Just vs        -- re-insert the tail
+    in
+    case fmap Seq.viewl mv of
+      Just (v Seq.:< _) -> return ( N nrs' (sz - fromIntegral key) , Just v  )
+      _                 -> return ( nrs,                             Nothing )
+
+
+-- | Add an entry to the nursery
 --
-{-# INLINE stash #-}
-stash :: forall m e proxy. RemoteMemory m
-      => proxy m
-      -> Int
-      -> NRS (RemotePtr m)
-      -> RemotePtr m e
-      -> IO ()
-stash _ !n !ref (castRemotePtr (Proxy :: Proxy m) -> ptr) = modifyMVar_ ref $ \(tbl,sz) -> do
-  mp  <- HT.lookup tbl n
-  case mp of
-    Nothing     -> HT.insert tbl n (FL.singleton () ptr)
-    Just xs     -> HT.insert tbl n (FL.cons () ptr xs)
-  return (tbl, sz + fromIntegral n)
-
-
--- Delete all entries from the nursery and free all associated device memory.
---
-flush :: (ptr Word8 -> IO ())
-      -> HashTable Int (FullList () (ptr Word8))
-      -> IO ()
-flush free !tbl =
-  let clean (!key,!val) = do
-        FL.mapM_ (const free) val
-        HT.delete tbl key
+{-# INLINEABLE insert #-}
+insert :: Int -> ptr Word8 -> Nursery ptr -> IO ()
+insert !key !val (Nursery !ref _) =
+  let
+      f Nothing   = Just (Seq.singleton val)
+      f (Just vs) = Just (vs Seq.|> val)
   in
-  message "flush nursery" >> HT.mapM_ clean tbl
+  modifyMVar_ ref $ \(N im sz) ->
+    return $! N (IM.alter f key im) (sz + fromIntegral key)
+
+
+-- | Delete all entries from the nursery
+--
+{-# INLINEABLE cleanup #-}
+cleanup :: (ptr Word8 -> IO ()) -> NRS ptr -> IO ()
+cleanup delete !ref = do
+  message "nursery cleanup"
+  modifyMVar_ ref $ \(N nrs _) -> do mapM_ (Seq.mapM delete) (IM.elems nrs)
+                                     return ( N IM.empty 0 )
+
 
--- The total size of all arrays stashed in the nursery.
+-- | The total number of bytes retained by the nursery
 --
+{-# INLINEABLE size #-}
 size :: Nursery ptr -> IO Int64
-size (Nursery ref _) = withMVar ref (return . snd)
+size (Nursery ref _) = withMVar ref $ \(N _ sz) -> return sz
 
 
 -- Debug

Data/Array/Accelerate/Array/Remote/Table.hs

@@ -37,14 +37,15 @@ module Data.Array.Accelerate.Array.Remote.Table (
 ) where
 
 import Control.Concurrent                                       ( yield )
-import Control.Concurrent.MVar                                  ( MVar, newMVar, withMVar, modifyMVar_, mkWeakMVar )
+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 )
 import Data.Maybe                                               ( isJust )
 import Data.Proxy
 import Data.Typeable                                            ( Typeable, gcast )
+import Data.Word
 import Foreign.Storable                                         ( sizeOf )
 import System.Mem                                               ( performGC )
 import System.Mem.Weak                                          ( Weak, deRefWeak )
@@ -82,7 +83,7 @@ type MT p               = MVar ( HashTable StableArray (RemoteArray p) )
 data MemoryTable p      = MemoryTable {-# UNPACK #-} !(MT p)
                                       {-# UNPACK #-} !(Weak (MT p))
                                       {-# UNPACK #-} !(Nursery p)
-                                      (forall a. p a -> IO ())
+                                      (p Word8 -> IO ())
 
 data RemoteArray p where
   RemoteArray :: Typeable e
@@ -173,12 +174,12 @@ malloc mt@(MemoryTable _ _ !nursery _) !ad !n = do
   message ("malloc: " ++ showBytes bytes)
   mp <-
     fmap (castRemotePtr (Proxy :: Proxy m))
-    <$> attempt "malloc/nursery" (liftIO $ N.malloc bytes nursery)
+    <$> attempt "malloc/nursery" (liftIO $ N.lookup bytes nursery)
         `orElse`
         attempt "malloc/new" (mallocRemote bytes)
         `orElse` do message "malloc/remote-malloc-failed (cleaning)"
                     clean mt
-                    liftIO $ N.malloc bytes nursery
+                    liftIO $ N.lookup bytes nursery
         `orElse` do message "malloc/remote-malloc-failed (purging)"
                     purge mt
                     mallocRemote bytes
@@ -218,6 +219,7 @@ free proxy mt !arr = do
   sa <- makeStableArray arr
   freeStable proxy mt sa
 
+
 -- | Deallocate the device array associated with the given StableArray. This
 -- is useful for other memory managers built on top of the memory table.
 --
@@ -227,13 +229,15 @@ freeStable
     -> MemoryTable (RemotePtr m)
     -> StableArray
     -> IO ()
-freeStable proxy (MemoryTable !ref _ (Nursery !nrs _) _) !sa = withMVar ref $ \mt -> do
-  mw <-  mt `HT.lookup` sa
-  case mw of
-    Nothing -> message ("free/already-removed: " ++ show sa)
-    Just (RemoteArray _ !p !bytes)  -> trace ("free/evict: " ++ show sa ++ " of " ++ showBytes bytes) $ do
-      N.stash proxy bytes nrs p
-      mt `HT.delete` sa
+freeStable proxy (MemoryTable !ref _ !nrs _) !sa =
+  withMVar ref $ \mt -> do
+    mw <- mt `HT.lookup` sa
+    case mw of
+      Nothing                        -> message ("free/already-removed: " ++ show sa)
+      Just (RemoteArray _ !p !bytes) -> do
+        message ("free/evict: " ++ show sa ++ " of " ++ showBytes bytes)
+        N.insert bytes (castRemotePtr proxy p) nrs
+        mt `HT.delete` sa
 
 
 -- Record an association between a host-side array and a new device memory area.
@@ -298,12 +302,15 @@ clean mt@(MemoryTable _ weak_ref nrs _) = management "clean" nrs . liftIO $ do
       alive <- isJust <$> deRefWeak w
       if alive then return rs else return (sa:rs)
 
--- |Call `free` on all arrays that are not currently associated with host-side
+
+-- | 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 _ _ nursery@(Nursery nrs _) release) = management "purge" nursery . liftIO $
-  modifyMVar_ nrs (\(tbl,_) -> N.flush release tbl >> return (tbl, 0))
+purge (MemoryTable _ _ nursery@(Nursery nrs _) release)
+  = management "purge" nursery
+  $ liftIO (N.cleanup release nrs)
+
 
 -- |Initiate garbage collection and `free` any remote arrays that no longer
 -- have matching host-side equivalents.