第53回　repaを使って行列演算を並列化する

-- | Parallel computation of array elements.
--
--   * The source array must have a delayed representation like `D`, `C` or `P`, 
--     and the result a manifest representation like `U` or `F`.
--
--   * If you want to copy data between manifest representations then use
--    `copyP` instead.
--
--   * If you want to convert a manifest array back to a delayed representation
--     then use `delay` instead.
--
computeP 
        :: (Shape sh, Fill r1 r2 sh e, Repr r2 e, Monad m)
        => Array r1 sh e -> m (Array r2 sh e)
computeP arr = now $ suspendedComputeP arr
{-# INLINE [4] computeP #-}


-- | Sequential computation of array elements.
computeS 
        :: Fill r1 r2 sh e
        => Array r1 sh e -> Array r2 sh e
computeS arr1
 = arr1 `deepSeqArray` 
   unsafePerformIO
 $ do   marr2    <- newMArr (size $ extent arr1) 
        fillS arr1 marr2
        unsafeFreezeMArr (extent arr1) marr2
{-# INLINE [4] computeS #-}

computeP 
        :: (Shape sh, Fill r1 r2 sh e, Repr r2 e, Monad m)
        => Array r1 sh e -> m (Array r2 sh e)
～ 略 ～

-- | Sequential computation of array elements.
computeS 
        :: Fill r1 r2 sh e
        => Array r1 sh e -> Array r2 sh e
～ 略 ～

import Prelude hiding (map)
import Data.Array.Repa

main = do
    let xs' = fromListUnboxed (ix3 2 3 4) [1..24]
    xs  <- computeP $ map (subtract 1) xs'
    ys  <- computeP $ map (subtract 1) (xs :: Array U DIM3 Double)
    zs  <- computeP $ map (/2) xs'
    print $ toList (ys :: Array U DIM3 Double)
    print $ toList (zs :: Array U DIM3 Double)

$ ./Test +RTS -N
[-1.0,0.0,1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0,11.0,12.0,13.0,14.0,15.0,16.0
,17.0,18.0,19.0,20.0,21.0,22.0]
[0.5,1.0,1.5,2.0,2.5,3.0,3.5,4.0,4.5,5.0,5.5,6.0,6.5,7.0,7.5,8.0,8.5,9.0,9.5,10.
0,10.5,11.0,11.5,12.0]

-- | Suspended parallel computation of array elements.
--
--   This version creates a thunk that will evaluate the array on demand.
--   If you force it when another parallel computation is already running
--   then you  will get a runtime warning and evaluation will be sequential. 
--   Use `deepSeqArray` and `now` to ensure that each array is evaluated
--   before proceeding to the next one. 
--  
--   If unsure then just use the monadic version `computeP`. This one ensures
--   that each array is fully evaluated before continuing.
--
suspendedComputeP 
        :: Fill r1 r2 sh e
        => Array r1 sh e -> Array r2 sh e
suspendedComputeP arr1
 = arr1 `deepSeqArray` 
   unsafePerformIO
 $ do   marr2    <- newMArr (size $ extent arr1) 
        fillP arr1 marr2
        unsafeFreezeMArr (extent arr1) marr2
{-# INLINE [4] suspendedComputeP #-}

-- | Sequential computation of array elements.
computeS 
        :: Fill r1 r2 sh e
        => Array r1 sh e -> Array r2 sh e
computeS arr1
 = arr1 `deepSeqArray` 
   unsafePerformIO
 $ do   marr2    <- newMArr (size $ extent arr1) 
        fillS arr1 marr2
        unsafeFreezeMArr (extent arr1) marr2
{-# INLINE [4] computeS #-}

-- Fillable -------------------------------------------------------------------
-- | Class of manifest array representations that can be filled in parallel 
--   and then frozen into immutable Repa arrays.
class Fillable r e where

 -- | Mutable version of the representation.
 data MArr r e

 -- | Allocate a new mutable array of the given size.
 newMArr          :: Int -> IO (MArr r e)

 -- | Write an element into the mutable array.
 unsafeWriteMArr  :: MArr r e -> Int -> e -> IO ()

 -- | Freeze the mutable array into an immutable Repa array.
 unsafeFreezeMArr :: sh  -> MArr r e -> IO (Array r sh e)

 -- | Ensure the strucure of a mutable array is fully evaluated.
 deepSeqMArr      :: MArr r e -> a -> a

 -- | Ensure the array is still live at this point.
 --   Needed when the mutable array is a ForeignPtr with a finalizer.
 touchMArr        :: MArr r e -> IO ()

-- Fill -----------------------------------------------------------------------
-- | Filling of unboxed vector arrays.
instance U.Unbox e => Fillable U e where
 data MArr U e 
  = UMArr (UM.IOVector e)

 newMArr n
  = liftM UMArr (UM.new n)
 {-# INLINE newMArr #-}

 unsafeWriteMArr (UMArr v) ix
  = UM.unsafeWrite v ix
 {-# INLINE unsafeWriteMArr #-}

 unsafeFreezeMArr sh (UMArr mvec)     
  = do  vec     <- U.unsafeFreeze mvec
        return  $  AUnboxed sh vec
 {-# INLINE unsafeFreezeMArr #-}

 deepSeqMArr (UMArr vec) x
  = vec `seq` x
 {-# INLINE deepSeqMArr #-}

 touchMArr _ 
  = return ()
 {-# INLINE touchMArr #-}

-- Fill -----------------------------------------------------------------------
-- | Compute all elements defined by an array and write them to a fillable
--   representation.
--  
--   Note that instances require that the source array to have a delayed
--   representation such as `D` or `C`. If you want to use a pre-existing
--   manifest array as the source then `delay` it first.
class (Shape sh, Repr r1 e, Fillable r2 e) => Fill r1 r2 sh e where
 -- | Fill an entire array sequentially.
 fillS          :: Array r1 sh e -> MArr r2 e -> IO ()

 -- | Fill an entire array in parallel.
 fillP          :: Array r1 sh e -> MArr r2 e -> IO ()

-- Fill -----------------------------------------------------------------------
-- | Compute all elements in an array.
instance (Fillable r2 e, Shape sh) => Fill D r2 sh e where
 fillP (ADelayed sh getElem) marr
  = marr `deepSeqMArr` 
    do  traceEventIO "Repa.fillP[Delayed]: start"
        fillChunkedP (size sh) (unsafeWriteMArr marr) (getElem . fromIndex sh) 
        touchMArr marr
        traceEventIO "Repa.fillP[Delayed]: end"
 {-# INLINE [4] fillP #-}

 fillS (ADelayed sh getElem) marr
  = marr `deepSeqMArr` 
    do  traceEventIO "Repa.fillS[Delayed]: start"
        fillChunkedS (size sh) (unsafeWriteMArr marr) (getElem . fromIndex sh)
        touchMArr marr
        traceEventIO "Repa.fillS[Delayed]: end"

 {-# INLINE [4] fillS #-}

-- | Monadic version of `deepSeqArray`. 
-- 
--   Forces an suspended array computation to be completed at this point
--   in a monadic computation.
--
-- @ do  let arr2 = suspendedComputeP arr1
--     ...
--     arr3 <- now $ arr2
--     ...
-- @
--
now     :: (Shape sh, Repr r e, Monad m)
        => Array r sh e -> m (Array r sh e)
now arr
 = do   arr `deepSeqArray` return ()
        return arr
{-# INLINE [4] now #-}

Prelude Data.Array.Repa> :t deepSeqArray
deepSeqArray :: (Shape sh, Repr r e) => Array r sh e -> b -> b

-- Conversions ----------------------------------------------------------------
-- | Sequential computation of array elements..
--
--   * This is an alias for `computeS` with a more specific type.
--
computeUnboxedS
        :: Fill r1 U sh e
        => Array r1 sh e -> Array U sh e
computeUnboxedS = computeS
{-# INLINE computeUnboxedS #-}


-- | Parallel computation of array elements.
--
--   * This is an alias for `computeP` with a more specific type.
--
computeUnboxedP
        :: (Fill r1 U sh e, Monad m, U.Unbox e)
        => Array r1 sh e -> m (Array U sh e)
computeUnboxedP = computeP
{-# INLINE computeUnboxedP #-}