Add combine2ByTagUP
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Mon, 9 Nov 2009 14:53:37 +0000 (14:53 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Mon, 9 Nov 2009 14:53:37 +0000 (14:53 +0000)
dph-prim-par/Data/Array/Parallel/Unlifted/Parallel.hs
dph-prim-par/Data/Array/Parallel/Unlifted/Parallel/Combinators.hs

index 6f4041b..e71333a 100644 (file)
@@ -18,7 +18,8 @@ module Data.Array.Parallel.Unlifted.Parallel (
 
   enumFromToUP, enumFromThenToUP, enumFromStepLenUP, enumFromStepLenEachUP,
 
-  mapUP, filterUP, packUP, combineUP, zipWithUP, foldUP, scanUP,
+  mapUP, filterUP, packUP, combineUP, combine2ByTagUP,
+  zipWithUP, foldUP, scanUP,
 
   andUP, sumUP,
   
index 49b3e58..0014050 100644 (file)
@@ -18,8 +18,8 @@
 #include "fusion-phases.h"
 
 module Data.Array.Parallel.Unlifted.Parallel.Combinators (
-  mapUP, filterUP, packUP, combineUP, zipWithUP, foldUP, fold1UP, foldl1UP,
-  scanUP
+  mapUP, filterUP, packUP, combineUP, combine2ByTagUP,
+  zipWithUP, foldUP, fold1UP, foldl1UP, scanUP
 ) where
 
 import Data.Array.Parallel.Base
@@ -62,6 +62,25 @@ combineUP flags !xs !ys = joinD theGang balanced
 
     go ((i :*: j) :*: (m :*: n)) bs = combineU bs (sliceU xs i m) (sliceU ys j n)
 
+combine2ByTagUP :: UA a => UArr Int -> UArr a -> UArr a -> UArr a
+{-# INLINE_UP combine2ByTagUP #-}
+combine2ByTagUP tags !xs !ys = joinD theGang balanced
+                             $ zipWithD theGang go (zipD is ns)
+                             $ splitD theGang balanced tags
+  where
+    ns = mapD theGang count
+       $ splitD theGang balanced tags
+
+    count bs = let ones = sumU bs
+               in (lengthU bs - ones) :*: ones
+
+    is = fstS $ scanD theGang add (0 :*: 0) ns
+
+    add (x1 :*: y1) (x2 :*: y2) = (x1+x2) :*: (y1+y2)
+
+    go ((i :*: j) :*: (m :*: n)) ts = combine2ByTagU ts (sliceU xs i m)
+                                                        (sliceU ys j n)
+
 
 zipWithUP :: (UA a, UA b, UA c) => (a -> b -> c) -> UArr a -> UArr b -> UArr c
 {-# INLINE zipWithUP #-}