Add zip3P, zipWith3P & unzip3P
authorManuel M T Chakravarty <chak@cse.unsw.edu.au>
Tue, 13 Sep 2011 11:11:43 +0000 (21:11 +1000)
committerManuel M T Chakravarty <chak@cse.unsw.edu.au>
Tue, 13 Sep 2011 11:11:43 +0000 (21:11 +1000)
dph-common/Data/Array/Parallel.hs
dph-common/Data/Array/Parallel/Lifted/Closure.hs
dph-common/Data/Array/Parallel/Lifted/Combinators.hs
dph-common/Data/Array/Parallel/PArray/PDataInstances.hs

index 932df3a..39f05cc 100644 (file)
@@ -36,7 +36,7 @@ module Data.Array.Parallel (
   (+:+), concatP,
   mapP, filterP, combineP,
   {- minimumP, maximumP, sumP, productP, -}  -- removed until we support type classes
-  zipP, unzipP, zipWithP,
+  zipP, zip3P, unzipP, unzip3P, zipWithP, zipWith3P,
   {- enumFromToP, enumFromThenToP, -}        -- removed until we support type classes
   bpermuteP, updateP, indexedP, sliceP,
   crossMapP,
@@ -135,16 +135,31 @@ zipP :: [:a:] -> [:b:] -> [:(a, b):]
 zipP !_ !_ = [::]
 {-# VECTORISE zipP = zipPA #-}
 
+zip3P :: [:a:] -> [:b:] -> [:c:] -> [:(a, b, c):]
+{-# NOINLINE zip3P #-}
+zip3P !_ !_ !_ = [::]
+{-# VECTORISE zip3P = zip3PA #-}
+
 unzipP :: [:(a, b):] -> ([:a:], [:b:])
 {-# NOINLINE unzipP #-}
 unzipP !_ = ([::], [::])
 {-# VECTORISE unzipP = unzipPA #-}
 
+unzip3P :: [:(a, b, c):] -> ([:a:], [:b:], [:c:])
+{-# NOINLINE unzip3P #-}
+unzip3P !_ = ([::], [::], [::])
+{-# VECTORISE unzip3P = unzip3PA #-}
+
 zipWithP :: (a -> b -> c) -> [:a:] -> [:b:] -> [:c:]
 {-# NOINLINE zipWithP #-}
 zipWithP !_ !_ !_ = [::]
 {-# VECTORISE zipWithP = zipWithPA #-}
 
+zipWith3P :: (a -> b -> c -> d) -> [:a:] -> [:b:] -> [:c:] -> [:d:]
+{-# NOINLINE zipWith3P #-}
+zipWith3P !_ !_ !_ !_ = [::]
+{-# VECTORISE zipWith3P = zipWith3PA #-}
+
 -- enumFromToP :: Enum a => a -> a -> [:a:]
 -- {-# NOINLINE enumFromToP #-}
 -- enumFromToP x y = [:x, y:]
index 6f6ec93..828ec7b 100644 (file)
@@ -4,7 +4,7 @@ module Data.Array.Parallel.Lifted.Closure (
   mkClosure, mkClosureP, ($:), ($:^),
   closure, liftedClosure, liftedApply,
 
-  closure1, closure2, closure3
+  closure1, closure2, closure3, closure4
 ) where
 import Data.Array.Parallel.PArray.PReprInstances ()
 import Data.Array.Parallel.PArray.PDataInstances
@@ -192,7 +192,6 @@ closure1 :: (a -> b) -> (PArray a -> PArray b) -> (a :-> b)
 {-# INLINE closure1 #-}
 closure1 fv fl = mkClosure (\_ -> fv) (\_ -> fl) ()
 
-
 -- | Arity-2 closures.
 closure2 :: PA a
          => (a -> b -> c)
@@ -205,7 +204,6 @@ closure2 fv fl = mkClosure fv_1 fl_1 ()
     fv_1 _ x  = mkClosure  fv fl x
     fl_1 _ xs = mkClosureP fv fl xs
 
-
 -- | Arity-3 closures.
 closure3 :: (PA a, PA b)
          => (a -> b -> c -> d)
@@ -224,3 +222,23 @@ closure3 fv fl = mkClosure fv_1 fl_1 ()
     fv_3 (x,y) z = fv x y z
     fl_3 ps zs = case unzipPA# ps of (xs,ys) -> fl xs ys zs
 
+-- | Arity-4 closures.
+closure4 :: (PA a, PA b, PA c)
+         => (a -> b -> c -> d -> e)
+         -> (PArray a -> PArray b -> PArray c -> PArray d -> PArray e)
+         -> (a :-> b :-> c :-> d :-> e)
+
+{-# INLINE closure4 #-}
+closure4 fv fl = mkClosure fv_1 fl_1 ()
+  where
+    fv_1 _  x  = mkClosure  fv_2 fl_2 x
+    fl_1 _  xs = mkClosureP fv_2 fl_2 xs
+
+    fv_2 x  y  = mkClosure  fv_3 fl_3 (x, y)
+    fl_2 xs ys = mkClosureP fv_3 fl_3 (zipPA# xs ys)
+
+    fv_3 (x, y)  z  = mkClosure  fv_4 fl_4 (x, y, z)
+    fl_3 xys     zs = case unzipPA# xys of (xs, ys) -> mkClosureP fv_4 fl_4 (zip3PA# xs ys zs)
+
+    fv_4 (x, y, z) v = fv x y z v
+    fl_4 ps vs = case unzip3PA# ps of (xs, ys, zs) -> fl xs ys zs vs
index aa99de1..419ae53 100644 (file)
@@ -23,7 +23,7 @@
 --
 module Data.Array.Parallel.Lifted.Combinators (
   lengthPA, replicatePA, singletonPA, mapPA, crossMapPA,
-  zipWithPA, zipPA, unzipPA, 
+  zipPA, zip3PA, zipWithPA, zipWith3PA, unzipPA, unzip3PA, 
   packPA, filterPA, combine2PA, indexPA, concatPA, appPA, enumFromToPA_Int,
   indexedPA, slicePA, updatePA, bpermutePA,
 
@@ -157,26 +157,43 @@ crossMapPA_l ass fs = copySegdPA# bss (zipPA# as' (concatPA# bss))
 
 
 -- zip ------------------------------------------------------------------------
--- | Takes two arrays and returns an array of corresponding pairs.
---   If one array is short, excess elements of the longer array are discarded.
-zipPA :: (PA a, PA b) => PArray a :-> PArray b :-> PArray (a,b)
+-- |Turns a tuple of arrays into an array of the corresponding tuples.
+--
+--  If one array is short, excess elements of the longer array are discarded.
+
+zipPA :: (PA a, PA b) => PArray a :-> PArray b :-> PArray (a, b)
 {-# INLINE zipPA #-}
 zipPA = closure2 zipPA_v zipPA_l
 
-zipPA_v :: (PA a, PA b) => PArray a -> PArray b -> PArray (a,b)
+zipPA_v :: (PA a, PA b) => PArray a -> PArray b -> PArray (a, b)
 {-# INLINE_PA zipPA_v #-}
 zipPA_v xs ys = zipPA# xs ys
 
 zipPA_l :: (PA a, PA b)
-        => PArray (PArray a) -> PArray (PArray b) -> PArray (PArray (a,b))
+        => PArray (PArray a) -> PArray (PArray b) -> PArray (PArray (a, b))
 {-# INLINE_PA zipPA_l #-}
 zipPA_l (PArray n# (PNested segd xs)) (PArray _ (PNested _ ys))
   = PArray n# (PNested segd (P_2 xs ys))
 
 
+zip3PA :: (PA a, PA b, PA c) => PArray a :-> PArray b :-> PArray c :-> PArray (a, b, c)
+{-# INLINE zip3PA #-}
+zip3PA = closure3 zip3PA_v zip3PA_l
+
+zip3PA_v :: (PA a, PA b, PA c) => PArray a -> PArray b -> PArray c -> PArray (a, b, c)
+{-# INLINE_PA zip3PA_v #-}
+zip3PA_v xs ys = zip3PA# xs ys
+
+zip3PA_l :: (PA a, PA b, PA c)
+         => PArray (PArray a) -> PArray (PArray b) -> PArray (PArray c) -> PArray (PArray (a, b, c))
+{-# INLINE_PA zip3PA_l #-}
+zip3PA_l (PArray n# (PNested segd xs)) (PArray _ (PNested _ ys)) (PArray _ (PNested _ zs))
+  = PArray n# (PNested segd (P_3 xs ys zs))
+
+
 -- zipWith --------------------------------------------------------------------
--- | zipWith generalises zip by zipping with the function given as the first
---   argument, instead of a tupling function.
+-- |Map a function over multiple arrays at once.
+
 zipWithPA :: (PA a, PA b, PA c)
           => (a :-> b :-> c) :-> PArray a :-> PArray b :-> PArray c
 {-# INLINE zipWithPA #-}
@@ -196,23 +213,57 @@ zipWithPA_l fs ass bss
       (replicatelPA# (segdPA# ass) fs $:^ concatPA# ass $:^ concatPA# bss)
 
 
+zipWith3PA :: (PA a, PA b, PA c, PA d)
+           => (a :-> b :-> c :-> d) :-> PArray a :-> PArray b :-> PArray c :-> PArray d
+{-# INLINE zipWith3PA #-}
+zipWith3PA = closure4 zipWith3PA_v zipWith3PA_l
+
+zipWith3PA_v :: (PA a, PA b, PA c, PA d)
+             => (a :-> b :-> c :-> d) -> PArray a -> PArray b -> PArray c -> PArray d
+{-# INLINE_PA zipWith3PA_v #-}
+zipWith3PA_v f as bs cs = replicatePA# (lengthPA# as) f $:^ as $:^ bs $:^ cs
+
+zipWith3PA_l :: (PA a, PA b, PA c, PA d)
+             => PArray (a :-> b :-> c :-> d) 
+             -> PArray (PArray a) -> PArray (PArray b) -> PArray (PArray c)
+             -> PArray (PArray d)
+{-# INLINE_PA zipWith3PA_l #-}
+zipWith3PA_l fs ass bss css
+  = copySegdPA# ass
+      (replicatelPA# (segdPA# ass) fs $:^ concatPA# ass $:^ concatPA# bss $:^ concatPA# css)
+
+
 -- unzip ----------------------------------------------------------------------
--- | Transform an array into an array of the first components,
---   and an array of the second components.
-unzipPA:: (PA a, PA b) => PArray (a, b) :-> (PArray a, PArray b)
+-- |Transform an array of tuples into a tuple of arrays.
+
+unzipPA :: (PA a, PA b) => PArray (a, b) :-> (PArray a, PArray b)
 {-# INLINE unzipPA #-}
 unzipPA = closure1 unzipPA_v unzipPA_l
 
-unzipPA_v:: (PA a, PA b) => PArray (a,b) -> (PArray a, PArray b)
+unzipPA_v :: (PA a, PA b) => PArray (a, b) -> (PArray a, PArray b)
 {-# INLINE_PA unzipPA_v #-}
 unzipPA_v abs' = unzipPA# abs'
 
-unzipPA_l:: (PA a, PA b) => PArray (PArray (a, b)) -> PArray (PArray a, PArray b)
+unzipPA_l :: (PA a, PA b) => PArray (PArray (a, b)) -> PArray (PArray a, PArray b)
 {-# INLINE_PA unzipPA_l #-}
 unzipPA_l xyss = zipPA# (copySegdPA# xyss xs) (copySegdPA# xyss ys)
   where
     (xs, ys) = unzipPA# (concatPA# xyss)
 
+unzip3PA :: (PA a, PA b, PA c) => PArray (a, b, c) :-> (PArray a, PArray b, PArray c)
+{-# INLINE unzip3PA #-}
+unzip3PA = closure1 unzip3PA_v unzip3PA_l
+
+unzip3PA_v :: (PA a, PA b, PA c) => PArray (a, b, c) -> (PArray a, PArray b, PArray c)
+{-# INLINE_PA unzip3PA_v #-}
+unzip3PA_v abs' = unzip3PA# abs'
+
+unzip3PA_l :: (PA a, PA b) => PArray (PArray (a, b, c)) -> PArray (PArray a, PArray b, PArray c)
+{-# INLINE_PA unzip3PA_l #-}
+unzip3PA_l xyzss = zip3PA# (copySegdPA# xyzss xs) (copySegdPA# xyzss ys) (copySegdPA# xyzss zs)
+  where
+    (xs, ys, zs) = unzip3PA# (concatPA# xyzss)
+
 
 -- packPA ---------------------------------------------------------------------
 -- | Select the elements of an array that have their tag set as True.
index 53cd78e..929c582 100644 (file)
@@ -11,7 +11,7 @@ module Data.Array.Parallel.PArray.PDataInstances(
   punit,
 
   -- * Operators on arrays of tuples
-  zipPA#,  unzipPA#, zip3PA#,
+  zipPA#,  unzipPA#, zip3PA#, unzip3PA#,
   
   -- * Operators on nested arrays
   segdPA#, concatPA#, segmentPA#, copySegdPA#
@@ -264,12 +264,16 @@ unzipPA# :: PArray (a, b) -> (PArray a, PArray b)
 unzipPA# (PArray n# (P_2 xs ys))
   = (PArray n# xs, PArray n# ys)
 
-
 zip3PA# :: PArray a -> PArray b -> PArray c -> PArray (a, b, c)
 {-# INLINE_PA zip3PA# #-}
 zip3PA# (PArray n# xs) (PArray _ ys) (PArray _ zs)
   = PArray n# (P_3 xs ys zs)
 
+unzip3PA# :: PArray (a, b, c) -> (PArray a, PArray b, PArray c)
+{-# INLINE_PA unzip3PA# #-}
+unzip3PA# (PArray n# (P_3 xs ys zs))
+  = (PArray n# xs, PArray n# ys, PArray n# zs)
+
 
 -- Sums -----------------------------------------------------------------------
 data instance PData (Sum2 a b)