Search code examples
haskellparallel-processingstate-monad

Parallelize computation of mutable vector in ST


How can computations done in ST be made to run in parallel?

I have a vector which needs to be filled in by random access, hence the use of ST, and the computation runs correctly single-threaded, but have been unable to figure out how to use more than one core.

Random access is needed because of the meaning of the indices into the vector. There are n things and every possible way of choosing among n things has an entry in the vector, as in the choice function. Each of these choices corresponds to a binary number (conceptually, a packed [Bool]) and these Int values are the indices. If there are n things, then the size of the vector is 2^n. The natural way the algorithm runs is for every entry corresponding to "n choose 1" to be filled in, then every entry for "n choose 2," etc. The entries corresponding to "n choose k" depends on the entries corresponding to "n choose (k-1)." The integers for the different choices do not occur in numerical order, and that's why random access is needed.

Here's a pointless (but slow) computation that follows the same pattern. The example function shows how I tried to break the computation up so that the bulk of the work is done in a pure world (no ST monad). In the code below, bogus is where most of the work is done, with the intent of calling that in parallel, but only one core is ever used.


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)


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]
    forM_ newEntries $ \e -> do
      let v = bogus p e
      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 (\k -> Vg.unsafeRead m k ) chs
  let e = map (\k -> toInteger k ) v
  return (sum e)

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

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

If anyone tests this, using more than 9 or 10 in show (example 9) will take much longer than you want to wait for such a pointless sequence of numbers.


Solution

  • Just do it in IO. If you need to use the result in pure code, then unsafePerformIO is available.

    The following version runs about 3-4 times faster with +RTS -N16 than +RTS -N1. My changes involved converting the ST vectors to IO, changing the forM_ to forConcurrently_, and adding a bang annotation to let !v = bogus ....

    Full code:

    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.Concurrent.Async
    import System.IO.Unsafe
    
    main :: IO ()
    main = do
      let m = unsafePerformIO (example 9)
      putStrLn $ show m
    
    example :: Int -> IO (Vb.Vector Int)
    example n = do
      m <- Vg.new (2^n)
    
      Vg.unsafeWrite m 0 (1)
    
      forM_ [1..n] $ \i -> do
        p <- prev m n (i-1)
        let newEntries = (choiceList n i) :: [Int]
        forConcurrently_ newEntries $ \e -> do
          let !v = bogus p e
          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.IOVector Int -> Int -> Int -> IO Integer
    prev m n 0 = return 1
    prev m n i = do
      let chs = choiceList n i
      v <- mapM (\k -> Vg.unsafeRead m k ) chs
      let e = map (\k -> toInteger k ) v
      return (sum e)
    
    bogus :: Integer -> Int -> Int
    bogus prior index = do
      let f = fac prior
      let g = (f^index) :: Integer
      let d = (map digitToInt (show g)) :: [Int]
      let a = fromIntegral (head d)^2
      a
    
    fac :: Integer -> Integer
    fac 0 = 1
    fac n = n * fac (n - 1)