module Data.Yarr.Fold (
Fold,
reduceL, reduceLeftM,
reduceR, reduceRightM,
runFold, runFoldP,
runFoldSlicesSeparate, runFoldSlicesSeparateP,
toList,
) where
import Prelude as P
import Control.Monad as M
import Data.List (groupBy)
import Data.Function (on)
import Data.Yarr.Base
import Data.Yarr.Shape as S
import Data.Yarr.Eval
import Data.Yarr.Convolution
import Data.Yarr.Utils.FixedVector as V hiding (toList)
import Data.Yarr.Utils.Fork
import Data.Yarr.Utils.Parallel
class Shape fsh => Reduce l ash fsh | l ash -> fsh where
getIndex :: USource r l ash a => UArray r l ash a -> (fsh -> IO a)
getSize :: USource r l ash a => UArray r l ash a -> fsh
#define SH_GET_INDEX(l,sh) \
instance Reduce l sh sh where { \
getIndex = index; \
getSize = extent; \
; \
; \
}
SH_GET_INDEX(SH, Dim1)
SH_GET_INDEX(SH, Dim2)
SH_GET_INDEX(SH, Dim3)
SH_GET_INDEX(CVL, Dim1)
SH_GET_INDEX(CVL, Dim2)
SH_GET_INDEX(CVL, Dim3)
#define LINEAR_GET_INDEX(l,sh) \
instance Reduce l sh Int where { \
getIndex = linearIndex; \
getSize = size . extent; \
; \
; \
}
LINEAR_GET_INDEX(L, Dim1)
LINEAR_GET_INDEX(L, Dim2)
LINEAR_GET_INDEX(L, Dim3)
type Fold sh a b =
IO b
-> (sh -> IO a)
-> sh
-> sh
-> IO b
reduceLeftM
:: Foldl sh a b
-> (b -> a -> IO b)
-> Fold sh a b
reduceLeftM foldl rf = foldl (\b _ a -> rf b a)
reduceL
:: Foldl sh a b
-> (b -> a -> b)
-> Fold sh a b
reduceL foldl rf = foldl (\b _ a -> return $ rf b a)
reduceRightM
:: Foldr sh a b
-> (a -> b -> IO b)
-> Fold sh a b
reduceRightM foldr rf = foldr (\_ a b -> rf a b)
reduceR
:: Foldr sh a b
-> (a -> b -> b)
-> Fold sh a b
reduceR foldr rf = foldr (\_ a b -> return $ rf a b)
runFold
:: (USource r l sh a, Reduce l sh fsh)
=> Fold fsh a b
-> IO b
-> UArray r l sh a
-> IO b
runFold fold mz arr = do
force arr
res <- fold mz (getIndex arr) zero (getSize arr)
touchArray arr
return res
runFoldP
:: (USource r l sh a, Reduce l sh fsh)
=> Threads
-> Fold fsh a b
-> IO b
-> (b -> b -> IO b)
-> UArray r l sh a
-> IO b
runFoldP threads fold mz join arr = do
force arr
ts <- threads
(r:rs) <- parallel ts $
makeFork ts zero (getSize arr) (fold mz (getIndex arr))
touchArray arr
M.foldM join r rs
runFoldSlicesSeparate
:: (UVecSource r slr l sh v e, Reduce l sh fsh)
=> Fold fsh e b
-> IO b
-> UArray r l sh (v e)
-> IO (VecList (Dim v) b)
runFoldSlicesSeparate fold mz arr =
V.mapM (\sl -> runFold fold mz sl) (slices arr)
runFoldSlicesSeparateP
:: (UVecSource r slr l sh v e, Reduce l sh fsh)
=> Threads
-> Fold fsh e b
-> IO b
-> (b -> b -> IO b)
-> UArray r l sh (v e)
-> IO (VecList (Dim v) b)
runFoldSlicesSeparateP threads fold mz join arr = do
force arr
ts <- threads
trs <- parallel ts $
makeForkSlicesOnce
ts
(V.replicate (zero, getSize arr))
(V.map (\sl -> fold mz (getIndex sl)) (slices arr))
touchArray arr
let rsBySlices = P.map (P.map snd) $ groupBy ((==) `on` fst) $ concat trs
rs <- M.mapM (\(r:rs) -> M.foldM join r rs) rsBySlices
return (VecList rs)
toList
:: (USource r l sh a, Reduce l sh fsh)
=> UArray r l sh a
-> IO [a]
toList = runFold (reduceR S.foldr (:)) (return [])