lfairy / ControlTest.hs

Created October 08, 2012

import Data.Vector.Unboxed as V
 
main :: IO ()
main = print $ sumVector ones
 
sumVector :: (Num a, Unbox a) => Vector a -> a
sumVector = V.foldl' (+) 0
 
{-# NOINLINE ones #-}
ones :: Vector Int
ones = V.replicate (800 * 1000) 1

-- | A simplified version of what happened at <http://www.reddit.com/r/programming/comments/1101bd>.
--
-- > ghc ParallelTest.hs -O2 -threaded
-- > ./ParallelTest +RTS -N
--
-- WARNING: ONLY RUN THIS PROGRAM ON A COMPUTER WITH AT LEAST 4 GB OF RAM!!!
 
import Control.Concurrent
import Control.Concurrent.MVar
import Data.Vector.Unboxed as V
import System.IO.Unsafe (unsafePerformIO)
 
main :: IO ()
main = print $ sumVector ones
 
-- | Add up all the numbers in a vector, in /parallel/!
sumVector :: (Num a, Unbox a) => Vector a -> a
sumVector vec = unsafePerformIO $ do
    dest <- newEmptyMVar
    run dest vec
    takeMVar dest
 
run :: (Num a, Unbox a) => MVar a -> Vector a -> IO ()
run dest vec
  | len == 0 = yield 0
  | len == 1 = yield $ V.head vec
  | otherwise = do
        let mid = len `div` 2
        lvar <- newEmptyMVar
        rvar <- newEmptyMVar
        forkIO $ run lvar (V.slice 0 mid vec)
        forkIO $ run rvar (V.slice mid (len - mid) vec)
        l <- takeMVar lvar
        r <- takeMVar rvar
        yield $! l + r
  where
    len = V.length vec
    yield = putMVar dest
 
{-# NOINLINE ones #-}
ones :: Vector Int
ones = V.replicate (800 * 1000) 1