如何在ST Monad中并行化可变向量的计算？

问题背景 #

我需要通过随机访问填充可变向量，因此采用了ST Monad，单线程运行正常，但无法利用多核资源。向量索引对应n个元素的所有选择组合（大小为2^n），算法按n选1、n选2的顺序填充条目：n选k的条目依赖n选k-1的结果，但同一k的条目之间无依赖，且选择对应的整数并非数值顺序，必须随机访问。我尝试将核心计算放到纯函数bogus中，但仍然只能单核心运行。

解决方案 #

同一批次（即同一k=i）的newEntries对应的计算是完全独立的——它们只依赖前一批的汇总值p和各自的索引e，互相之间没有依赖关系。因此我们可以把这部分纯计算逻辑抽出来，使用并行策略提前计算好所有结果，再批量写入ST向量。

具体实现需要用到Control.Parallel.Strategies中的parMap来并行处理每个bogus调用，再将结果写入可变向量。

修改后的代码 #

import qualified Data.Vector as Vb
import qualified Data.Vector.Mutable as Vm
import qualified Data.Vector.Generic.Mutable as Vg
import qualified Data.Vector.Generic as Gg
import Control.Monad.ST as ST ( ST, runST )
import Data.Foldable(forM_)
import Data.Char(digitToInt)
import Control.Parallel.Strategies (parMap, rseq) -- 新增并行相关导入

main :: IO ()
main = do
  putStrLn $ show (example 9)

example :: Int -> Vb.Vector Int
example n = runST $ do
  m <- Vg.new (2^n) :: ST s (Vm.STVector s Int)
  
  Vg.unsafeWrite m 0 1
  
  forM_ [1..n] $ \i -> do
    p <- prev m n (i-1)
    let newEntries = choiceList n i :: [Int]
    -- 并行计算所有结果，rseq确保每个计算完全求值后再收集
    let results = parMap rseq (\e -> (e, bogus p e)) newEntries
    -- 批量写入ST向量
    forM_ results $ \(e, v) -> Vg.unsafeWrite m e v
  
  Gg.unsafeFreeze m

choiceList :: Int -> Int -> [Int]
choiceList _ 0 = [0]
choiceList n 1 = [ 2^k | k <- [0..(n-1) ] ]
choiceList n k 
  | n == k = [2^n - 1]
  | otherwise = choiceList (n-1) k ++ map ((2^(n-1)) +) (choiceList (n-1) (k-1))

prev :: Vm.STVector s Int -> Int -> Int -> ST s Integer
prev m n 0 = return 1
prev m n i = do
  let chs = choiceList n i
  v <- mapM (Vg.unsafeRead m) chs
  return $ sum (map toInteger v)

bogus :: Integer -> Int -> Int
bogus prior index = 
  let f = fac prior
      g = f^index :: Integer
      d = map digitToInt (show g) :: [Int]
      a = fromIntegral (head d)^2
  in a

fac :: Integer -> Integer
fac 0 = 1
fac n = n * fac (n - 1)

关键说明 #

• parMap rseq会并行计算列表中每个元素的bogus p e值，rseq保证每个计算结果完全求值后再被收集，避免惰性求值带来的并行失效。
• 因为bogus是纯函数，同一批次的计算互相独立，并行执行不会产生竞态条件。
• ST Monad内部的向量写入仍然是单线程的，但这部分操作本身开销远小于bogus中的大整数运算（阶乘、幂运算），所以并行计算核心逻辑能有效利用多核资源。

注意事项 #

• 编译时需要加上-threaded选项，运行时需指定+RTS -N<核数>（比如+RTS -N4）来启用并行。
• 当n超过9或10时，2^n会快速增长，同时bogus中的大整数运算开销会急剧上升，即使并行也会耗时较久，这是算法本身的复杂度问题，而非并行实现的问题。

内容的提问来源于stack exchange，提问作者Randall Fairman