Repa 数组上的并行 mapM

在我最近的 工作与 Gibbs sampling，我一直在大量使用的 RVar，在我看来，提供了一个近乎理想的接口随机数生成。遗憾的是，由于无法在地图中使用单子动作，我无法使用 Repa。

虽然一元映射通常不能并行化，但在我看来，RVar至少是一个单元映射可以安全并行化的例子(至少在原则上是这样; 我对 RVar的内部工作原理并不十分熟悉)。也就是说，我想写下这样的东西,

drawClass :: Sample -> RVar Class
drawClass = ...


drawClasses :: Array U DIM1 Sample -> RVar (Array U DIM1 Class)
drawClasses samples = A.mapM drawClass samples

A.mapM看起来像是,

mapM :: ParallelMonad m => (a -> m b) -> Array r sh a -> m (Array r sh b)

显然，这将如何工作取决于 RVar及其底层 RandomSource的实现，原则上，人们会认为这将涉及为每个产生并照常进行的线程绘制一个新的随机种子。

凭直觉，似乎同样的想法可以推广到其他一些单子。

因此，我的问题是: 能否构建一个单子类 ParallelMonad，其效果可以安全地并行化(可能至少由 RVar居住) ？

它看起来像什么？还有哪些单子可能存在于这个类中？其他人是否考虑过这种方法在 Repa 的可能性？

最后，如果这种并行一元动作的概念不能被推广，那么在 RVar的特定情况下(在这种情况下它会非常有用) ，有人认为有什么好的方法可以实现这种功能吗？为了并行性而放弃 RVar是一个非常困难的权衡。

2441 次浏览
小开

It's probably not a good idea to do this due to inherently sequential nature of PRNGs. Instead, you might want to transition your code as follows:

  1. Declare an IO function (main, or what have you).
  2. Read as many random numbers as you need.
  3. Pass the (now pure) numbers onto your repa functions.
小开

It's been 7 years since this question has been asked, and it still seems like no-one came up with a good solution to this problem. Repa doesn't have a mapM/traverse like function, even one that could run without parallelization. Moreover, considering the amount of progress there was in the last few years it seems unlikely that it will happen either.

Because of stale state of many array libraries in Haskell and my overall dissatisfaction with their feature sets I've put forth a couple of years of work into an array library massiv, which borrows some concepts from Repa, but takes it to a completely different level. Enough with the intro.

Prior to today, there was three monadic map like functions in massiv (not counting the synonym like functions: imapM, forM et al.):

  • mapM - the usual mapping in an arbitrary Monad. Not parallelizable for obvious reasons and is also a bit slow (along the lines of usual mapM over a list slow)
  • traversePrim - here we are restricted to PrimMonad, which is significantly faster than mapM, but the reason for this is not important for this discussion.
  • mapIO - this one, as name suggests, is restricted to IO (or rather MonadUnliftIO, but that is irrelevant). Because we are in IO we can automatically split array in as many chunks as there are cores and use separate worker threads to map the IO action over each element in those chunks. Unlike pure fmap, which is also parallelizable, we have to be in IO here because of non-determinism of scheduling combined with side effects of our mapping action.

So, once I read this question, I thought to myself that the problem is practically solved in massiv, but no so fast. Random number generators, such as in mwc-random and others in random-fu can't use the same generator across many threads. Which means, that the only piece of the puzzle I was missing was: "drawing a new random seed for each thread spawned and proceeding as usual". In other words, I needed two things:

  • A function that would initialize as many generators as there gonna be worker threads
  • and an abstraction that would seamlessly give the correct generator to the mapping function depending on which thread that the action is running in.

So that is exactly what I did.

First I will give examples using the specially crafted randomArrayWS and initWorkerStates functions, as they are more relevant to the question and later move to the more general monadic map. Here are their type signatures:

randomArrayWS ::
(Mutable r ix e, MonadUnliftIO m, PrimMonad m)
=> WorkerStates g -- ^ Use `initWorkerStates` to initialize you per thread generators
-> Sz ix -- ^ Resulting size of the array
-> (g -> m e) -- ^ Generate the value using the per thread generator.
-> m (Array r ix e)

initWorkerStates :: MonadIO m => Comp -> (WorkerId -> m s) -> m (WorkerStates s)

For those who are not familiar with massiv, the Comp argument is a computation strategy to use, notable constructors are:

  • Seq - run computation sequentially, without forking any threads
  • Par - spin up as many threads as there are capabilities and use those to do the work.

I'll use mwc-random package as an example initially and later move to RVarT:

λ> import Data.Massiv.Array
λ> import System.Random.MWC (createSystemRandom, uniformR)
λ> import System.Random.MWC.Distributions (standard)
λ> gens <- initWorkerStates Par (\_ -> createSystemRandom)

Above we initialized a separate generator per thread using system randomness, but we could have just as well used a unique per thread seed by deriving it from the WorkerId argument, which is a mere Int index of the worker. And now we can use those generators to create an array with random values:

λ> randomArrayWS gens (Sz2 2 3) standard :: IO (Array P Ix2 Double)
Array P Par (Sz (2 :. 3))
[ [ -0.9066144845415213, 0.5264323240310042, -1.320943607597422 ]
, [ -0.6837929005619592, -0.3041255565826211, 6.53353089112833e-2 ]
]

By using Par strategy the scheduler library will split evenly the work of generation among available workers and each worker will use it's own generator, thus making it thread safe. Nothing prevents us from reusing the same WorkerStates arbitrary number of times as long as it is not done concurrently, which otherwise would result in an exception:

λ> randomArrayWS gens (Sz1 10) (uniformR (0, 9)) :: IO (Array P Ix1 Int)
Array P Par (Sz1 10)
[ 3, 6, 1, 2, 1, 7, 6, 0, 8, 8 ]

Now putting mwc-random to the side we can reuse the same concept for other possible uses cases by using functions like generateArrayWS:

generateArrayWS ::
(Mutable r ix e, MonadUnliftIO m, PrimMonad m)
=> WorkerStates s
-> Sz ix --  ^ size of new array
-> (ix -> s -> m e) -- ^ element generating action
-> m (Array r ix e)

and mapWS:

mapWS ::
(Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad m)
=> WorkerStates s
-> (a -> s -> m b) -- ^ Mapping action
-> Array r' ix a -- ^ Source array
-> m (Array r ix b)

Here is the promised example on how to use this functionality with rvar, random-fu and mersenne-random-pure64 libraries. We could have used randomArrayWS here as well, but for the sake of example let's say we already have an array with different RVarTs, in which case we need a mapWS:

λ> import Data.Massiv.Array
λ> import Control.Scheduler (WorkerId(..), initWorkerStates)
λ> import Data.IORef
λ> import System.Random.Mersenne.Pure64 as MT
λ> import Data.RVar as RVar
λ> import Data.Random as Fu
λ> rvarArray = makeArrayR D Par (Sz2 3 9) (\ (i :. j) -> Fu.uniformT i j)
λ> mtState <- initWorkerStates Par (newIORef . MT.pureMT . fromIntegral . getWorkerId)
λ> mapWS mtState RVar.runRVarT rvarArray :: IO (Array P Ix2 Int)
Array P Par (Sz (3 :. 9))
[ [ 0, 1, 2, 2, 2, 4, 5, 0, 3 ]
, [ 1, 1, 1, 2, 3, 2, 6, 6, 2 ]
, [ 0, 1, 2, 3, 4, 4, 6, 7, 7 ]
]

It is important to note, that despite the fact that pure implementation of Mersenne Twister is being used in the above example, we cannot escape the IO. This is because of the non-deterministic scheduling, which means that we never know which one of the workers will be handling which chunk of the array and consequently which generator will be used for which part of the array. On the up side, if the generator is pure and splittable, such as splitmix, then we can use the pure, deterministic and parallelizable generation function: randomArray, but that is already a separate story.


提交答案