Vector parallelization

Main.hs

module Main where

import GHC.Conc (numCapabilities)
import Control.Parallel.Strategies
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G

autoParVector :: (G.Vector v a, NFData (v a), NFData a) => Strategy (v a)
autoParVector = autoChunkStrategy id

autoParFilter :: (G.Vector v a, NFData (v a), NFData a) => (a -> Bool) -> Strategy (v a)
autoParFilter func = autoChunkStrategy (G.filter func)

parVector :: (NFData (v a), G.Vector v a, NFData a) => Int -> Strategy (v a)
parVector = parChunkStrategy id

parFilter :: (G.Vector v a, NFData (v a), NFData a) => Int -> (a -> Bool) -> Strategy (v a)
parFilter minChunk func = parChunkStrategy (G.filter func) minChunk

autoChunkStrategy :: (G.Vector v a, NFData (v a), NFData a) => (v a -> v a) -> Strategy (v a)
autoChunkStrategy func vec
  | numCapabilities > 1 = parChunkStrategy func (G.length vec `div` (10 * numCapabilities)) vec
  | otherwise           = rseq (func vec)

parChunkStrategy :: (G.Vector v a, NFData (v a), NFData a) => (v a -> v a) -> Int -> Strategy (v a)
parChunkStrategy func minChunk = fmap G.concat . innerPar
  where
    innerPar vec
      | vLen > minChunk = let
        half = vLen `div` 2
        v1 = G.unsafeSlice 0 half vec
        v2 = G.unsafeSlice half (vLen - half) vec
        in do
          a <- innerPar v1
          b <- innerPar v2
          return (a ++ b)
      | otherwise = do
          e <- rparWith rdeepseq (func vec)
          return [e]
      where
        vLen = G.length vec

fac :: Int -> Int
fac n
  | n < 2     = 1
  | otherwise = n * (fac $ n - 1)

main :: IO ()
main = do
  let result = (V.enumFromN 50000 50500) `using` autoParFilter (even . fac) 
  putStrLn . show $ G.sum result

  let resultMap = V.map fac (V.enumFromN 50000 50500) `using` autoParVector 
  putStrLn . show $ G.sum resultMap