<commit_before>{-# LANGUAGE CPP #-}
module Concurrent.Capability.LocalArray
  ( LocalArray
  , newLocalArray
  , readLocalArray
  , writeLocalArray
  , atomicModifyLocalArray
  , atomicModifyLocalArray'
  , modifyLocalArray
  , modifyLocalArray'
  , fetchModifyLocalArray
  , fetchModifyLocalArray'
  ) where

import Concurrent.Primitive
import Concurrent.Thread
import Control.Concurrent
import Control.Monad.Primitive
import Data.Bits
import Data.Primitive
import Foreign.Storable as S

-- | Capability-Local variables with cache-line spacing
--
-- You should really only access these from code that is pinned to a given capability, or when you just don't care which bin the result falls into.
data LocalArray a = LocalArray
  {-# UNPACK #-} !Int -- logical array size
  {-# UNPACK #-} !Int -- cache-line-expanded array size per capability
  {-# UNPACK #-} !(MutableArray RealWorld a) -- actual array

instance Eq (LocalArray a) where
  LocalArray _ _ x == LocalArray _ _ y = sameMutableArray x y

pointersPerCacheLine :: Int
pointersPerCacheLine
  | S.sizeOf (undefined :: Int) == 4 = 16
  | S.sizeOf (undefined :: Int) == 8 = 8
  | otherwise = error "unknown number of pointers per cache line"

logPointersPerCacheLine :: Int
logPointersPerCacheLine
  | S.sizeOf (undefined :: Int) == 4 = 4
  | S.sizeOf (undefined :: Int) == 8 = 3
  | otherwise = error "unknown number of pointers per cache line"

-- | Calling 'Control.Concurrent.setNumCapabilities' after this is built will cause you to crash when accessing it
-- and break invariants.
newLocalArray :: MonadPrimIO m => Int -> a -> m (LocalArray a)
newLocalArray n a = primIO $ do
  cs <- getNumCapabilities
  let up = unsafeShiftR (n + pointersPerCacheLine - 1) logPointersPerCacheLine
  LocalArray n up <$> newArray (up*(cs-1)+n) a

readLocalArray :: MonadPrimIO m => LocalArray a -> Int -> m a
readLocalArray (LocalArray _ o arr) i = do
  j <- currentCapability
  readArray arr ((o*j)+i)

writeLocalArray :: MonadPrimIO m => LocalArray a -> Int -> a -> m ()
writeLocalArray (LocalArray _ o arr) i a = do
  j <- currentCapability
  writeArray arr ((o*j)+i) a

-- | This is safely atomic, despite the lack of CAS.
atomicModifyLocalArray :: MonadPrimIO m => LocalArray a -> Int -> (a -> (a, b)) -> m b
atomicModifyLocalArray (LocalArray _ o arr) i f = do
  j <- currentCapability
  localAtomicModifyArray arr ((o*j)+i) f

-- | This is safely atomic, despite the lack of CAS.
atomicModifyLocalArray' :: MonadPrimIO m => LocalArray a -> Int -> (a -> (a, b)) -> m b
atomicModifyLocalArray' (LocalArray _ o arr) i f = do
  j <- currentCapability
  localAtomicModifyArray' arr ((o*j)+i) f

-- | This is safely atomic, despite the lack of CAS.
modifyLocalArray :: MonadPrimIO m => LocalArray a -> Int -> (a -> a) -> m a
modifyLocalArray (LocalArray _ o arr) i f = do
  j <- currentCapability
  localModifyArray arr ((o*j)+i) f

-- | This is safely atomic, despite the lack of CAS.
modifyLocalArray' :: MonadPrimIO m => LocalArray a -> Int -> (a -> a) -> m a
modifyLocalArray' (LocalArray _ o arr) i f = do
  j <- currentCapability
  localModifyArray' arr ((o*j)+i) f

-- | This is safely atomic, despite the lack of CAS.
fetchModifyLocalArray :: MonadPrimIO m => LocalArray a -> Int -> (a -> a) -> m a
fetchModifyLocalArray (LocalArray _ o arr) i f = do
  j <- currentCapability
  localFetchModifyArray arr ((o*j)+i) f

-- | This is safely atomic, despite the lack of CAS.
fetchModifyLocalArray' :: MonadPrimIO m => LocalArray a -> Int -> (a -> a) -> m a
fetchModifyLocalArray' (LocalArray _ o arr) i f = do
  j <- currentCapability
  localFetchModifyArray' arr ((o*j)+i) f<commit_msg>Switch LocalArray over to using a SmallMutableArray to avoid false cache-line sharing problems caused by the card marking table<commit_after>{-# LANGUAGE CPP #-}
module Concurrent.Capability.LocalArray
  ( LocalArray
  , newLocalArray
  , readLocalArray
  , writeLocalArray
  , atomicModifyLocalArray
  , atomicModifyLocalArray'
  , modifyLocalArray
  , modifyLocalArray'
  , fetchModifyLocalArray
  , fetchModifyLocalArray'
  ) where

import Concurrent.Primitive
import Concurrent.Thread
import Control.Concurrent
import Control.Monad.Primitive
import Data.Bits
import Foreign.Storable as S

-- | Capability-Local variables with cache-line spacing
--
-- You should really only access these from code that is pinned to a given capability, or when you just don't care which bin the result falls into.
data LocalArray a = LocalArray
  {-# UNPACK #-} !Int -- logical array size
  {-# UNPACK #-} !Int -- cache-line-expanded array size per capability
  {-# UNPACK #-} !(SmallMutableArray RealWorld a) -- actual array

instance Eq (LocalArray a) where
  LocalArray _ _ x == LocalArray _ _ y = sameSmallMutableArray x y

pointersPerCacheLine :: Int
pointersPerCacheLine
  | S.sizeOf (undefined :: Int) == 4 = 16
  | S.sizeOf (undefined :: Int) == 8 = 8
  | otherwise = error "unknown number of pointers per cache line"

logPointersPerCacheLine :: Int
logPointersPerCacheLine
  | S.sizeOf (undefined :: Int) == 4 = 4
  | S.sizeOf (undefined :: Int) == 8 = 3
  | otherwise = error "unknown number of pointers per cache line"

-- | Calling 'Control.Concurrent.setNumCapabilities' after this is built will cause you to crash when accessing it
-- and break invariants.
newLocalArray :: MonadPrimIO m => Int -> a -> m (LocalArray a)
newLocalArray n a = primIO $ do
  cs <- getNumCapabilities
  let up = unsafeShiftR (n + pointersPerCacheLine - 1) logPointersPerCacheLine
  LocalArray n up <$> newSmallArray (up*(cs-1)+n) a

readLocalArray :: MonadPrimIO m => LocalArray a -> Int -> m a
readLocalArray (LocalArray _ o arr) i = do
  j <- currentCapability
  readSmallArray arr ((o*j)+i)

writeLocalArray :: MonadPrimIO m => LocalArray a -> Int -> a -> m ()
writeLocalArray (LocalArray _ o arr) i a = do
  j <- currentCapability
  writeSmallArray arr ((o*j)+i) a

-- | This is safely atomic, despite the lack of CAS.
atomicModifyLocalArray :: MonadPrimIO m => LocalArray a -> Int -> (a -> (a, b)) -> m b
atomicModifyLocalArray (LocalArray _ o arr) i f = do
  j <- currentCapability
  localAtomicModifySmallArray arr ((o*j)+i) f

-- | This is safely atomic, despite the lack of CAS.
atomicModifyLocalArray' :: MonadPrimIO m => LocalArray a -> Int -> (a -> (a, b)) -> m b
atomicModifyLocalArray' (LocalArray _ o arr) i f = do
  j <- currentCapability
  localAtomicModifySmallArray' arr ((o*j)+i) f

-- | This is safely atomic, despite the lack of CAS.
modifyLocalArray :: MonadPrimIO m => LocalArray a -> Int -> (a -> a) -> m a
modifyLocalArray (LocalArray _ o arr) i f = do
  j <- currentCapability
  localModifySmallArray arr ((o*j)+i) f

-- | This is safely atomic, despite the lack of CAS.
modifyLocalArray' :: MonadPrimIO m => LocalArray a -> Int -> (a -> a) -> m a
modifyLocalArray' (LocalArray _ o arr) i f = do
  j <- currentCapability
  localModifySmallArray' arr ((o*j)+i) f

-- | This is safely atomic, despite the lack of CAS.
fetchModifyLocalArray :: MonadPrimIO m => LocalArray a -> Int -> (a -> a) -> m a
fetchModifyLocalArray (LocalArray _ o arr) i f = do
  j <- currentCapability
  localFetchModifySmallArray arr ((o*j)+i) f

-- | This is safely atomic, despite the lack of CAS.
fetchModifyLocalArray' :: MonadPrimIO m => LocalArray a -> Int -> (a -> a) -> m a
fetchModifyLocalArray' (LocalArray _ o arr) i f = do
  j <- currentCapability
  localFetchModifySmallArray' arr ((o*j)+i) f
{'fraction_non_alphanumeric': 0.10944871276452352, 'fraction_numerical': 0.002965359212831918, 'mean_word_length': 3.7655748233782917, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 85, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5203096', 'n_tokens_mistral': 2412, 'n_tokens_neox': 2332, 'n_words': 857}