From 66dcf1fd28a17d100d2ef4e2914ce8c5d0a6c14c Mon Sep 17 00:00:00 2001 From: Duncan Coutts Date: Thu, 8 Aug 2024 11:54:05 +0100 Subject: [PATCH 1/2] SPECIALISE functions in Data.Arena They're overloaded on MonadPrim m => and are called in the hot path of lookup, where using the overloaded version was generating a lot of allocations. The alloc per key in the benchIndexSearches goes from 1369 to 88.9. --- src/Data/Arena.hs | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/src/Data/Arena.hs b/src/Data/Arena.hs index 1826ef38e..0ededd01c 100644 --- a/src/Data/Arena.hs +++ b/src/Data/Arena.hs @@ -20,6 +20,7 @@ module Data.Arena ( import Control.DeepSeq (NFData (..)) import Control.Exception (assert) import Control.Monad.Primitive +import Control.Monad.ST (ST) import Data.Bits (complement, popCount, (.&.)) import Data.Primitive.ByteArray import Data.Primitive.MutVar @@ -59,12 +60,19 @@ type Alignment = Int blockSize :: Int blockSize = 0x100000 +{-# SPECIALIZE + newBlock :: ST s (Block s) + #-} +{-# SPECIALIZE + newBlock :: IO (Block RealWorld) + #-} newBlock :: PrimMonad m => m (Block (PrimState m)) newBlock = do off <- newPrimVar 0 mba <- newAlignedPinnedByteArray blockSize 4096 return (Block off mba) +{-# INLINE withArena #-} withArena :: PrimMonad m => ArenaManager (PrimState m) -> (Arena (PrimState m) -> m a) -> m a withArena am f = do a <- newArena am @@ -72,6 +80,12 @@ withArena am f = do closeArena am a pure x +{-# SPECIALIZE + newArena :: ArenaManager s -> ST s (Arena s) + #-} +{-# SPECIALIZE + newArena :: ArenaManager RealWorld -> IO (Arena RealWorld) + #-} newArena :: PrimMonad m => ArenaManager (PrimState m) -> m (Arena (PrimState m)) newArena (ArenaManager arenas) = do marena <- atomicModifyMutVar' arenas $ \case @@ -86,6 +100,12 @@ newArena (ArenaManager arenas) = do full <- newMutVar [] return Arena {..} +{-# SPECIALIZE + closeArena :: ArenaManager s -> Arena s -> ST s () + #-} +{-# SPECIALIZE + closeArena :: ArenaManager RealWorld -> Arena RealWorld -> IO () + #-} closeArena :: PrimMonad m => ArenaManager (PrimState m) -> Arena (PrimState m) -> m () closeArena (ArenaManager arenas) arena = do scrambleArena arena @@ -112,6 +132,12 @@ scrambleBlock (Block _ mba) = do setByteArray mba 0 size (0x77 :: Word8) #endif +{-# SPECIALIZE + resetArena :: Arena s -> ST s () + #-} +{-# SPECIALIZE + resetArena :: Arena RealWorld -> IO () + #-} -- | Reset arena, i.e. return used blocks to free list. resetArena :: PrimMonad m => Arena (PrimState m) -> m () resetArena Arena {..} = do @@ -135,6 +161,9 @@ withUnmanagedArena k = do mgr <- newArenaManager withArena mgr k +{-# SPECIALIZE + allocateFromArena :: Arena s -> Size -> Alignment -> ST s (Offset, MutableByteArray s) + #-} -- | Allocate a slice of mutable byte array from the arena. allocateFromArena :: PrimMonad m => Arena (PrimState m)-> Size -> Alignment -> m (Offset, MutableByteArray (PrimState m)) allocateFromArena !arena !size !alignment = @@ -142,6 +171,9 @@ allocateFromArena !arena !size !alignment = assert (size <= blockSize) $ -- not too large allocations allocateFromArena' arena size alignment +{-# SPECIALIZE + allocateFromArena' :: Arena s -> Size -> Alignment -> ST s (Offset, MutableByteArray s) + #-} -- TODO!? this is not async exception safe allocateFromArena' :: PrimMonad m => Arena (PrimState m)-> Size -> Alignment -> m (Offset, MutableByteArray (PrimState m)) allocateFromArena' arena@Arena { .. } !size !alignment = do @@ -173,6 +205,7 @@ allocateFromArena' arena@Arena { .. } !size !alignment = do -- * go again allocateFromArena' arena size alignment +{-# SPECIALIZE newBlockWithFree :: MutVar s [Block s] -> ST s (Block s) #-} -- | Allocate new block, possibly taking it from a free list newBlockWithFree :: PrimMonad m => MutVar (PrimState m) [Block (PrimState m)] -> m (Block (PrimState m)) newBlockWithFree free = do From 74ba6c7b5169f9baa10ab0e63f638fffd0967a33 Mon Sep 17 00:00:00 2001 From: Duncan Coutts Date: Thu, 8 Aug 2024 11:56:20 +0100 Subject: [PATCH 2/2] Reduce indirections in IndexCompact by using UNPACK Strict fields are not auomatically unpacked when they're bigger than a word. So things like Vector etc are not unpacked automagically. --- src/Database/LSMTree/Internal/IndexCompact.hs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/Database/LSMTree/Internal/IndexCompact.hs b/src/Database/LSMTree/Internal/IndexCompact.hs index f827d5d42..186181032 100644 --- a/src/Database/LSMTree/Internal/IndexCompact.hs +++ b/src/Database/LSMTree/Internal/IndexCompact.hs @@ -360,11 +360,11 @@ import Database.LSMTree.Internal.Vector data IndexCompact = IndexCompact { -- | \(P\): Maps a page @i@ to the 64-bit slice of primary bits of its -- minimum key. - icPrimary :: !(VU.Vector Word64) + icPrimary :: {-# UNPACK #-} !(VU.Vector Word64) -- | \(C\): A clash on page @i@ means that the primary bits of the minimum -- key on that page aren't sufficient to decide whether a search for a key -- should continue left or right of the page. - , icClashes :: !(VU.Vector Bit) + , icClashes :: {-# UNPACK #-} !(VU.Vector Bit) -- | \(TB\): Maps a full minimum key to the page @i@ that contains it, but -- only if there is a clash on page @i@. , icTieBreaker :: !(Map (Unsliced SerialisedKey) PageNo) @@ -372,7 +372,7 @@ data IndexCompact = IndexCompact { -- the larger-than-page value, the first page will map to 'False', and the -- remainder of the pages will be set to 'True'. Regular pages default to -- 'False'. - , icLargerThanPage :: !(VU.Vector Bit) + , icLargerThanPage :: {-# UNPACK #-} !(VU.Vector Bit) } deriving stock (Show, Eq)