dph-lifted-vsegd: make replicatesPR take a Segd
authorBen Lippmeier <benl@ouroborus.net>
Thu, 13 Oct 2011 03:12:50 +0000 (14:12 +1100)
committerBen Lippmeier <benl@ouroborus.net>
Thu, 13 Oct 2011 03:12:50 +0000 (14:12 +1100)
We'll want to check whether there are zero segment lengths to use updateVSegsReachable

dph-lifted-vseg/Data/Array/Parallel/Lifted/Combinators.hs
dph-lifted-vseg/Data/Array/Parallel/PArray.hs
dph-lifted-vseg/Data/Array/Parallel/PArray/PData/Base.hs
dph-lifted-vseg/Data/Array/Parallel/PArray/PData/Double.hs
dph-lifted-vseg/Data/Array/Parallel/PArray/PData/Int.hs
dph-lifted-vseg/Data/Array/Parallel/PArray/PData/Nested.hs
dph-lifted-vseg/Data/Array/Parallel/PArray/PData/Unit.hs

index 99f81e8..683b21a 100644 (file)
@@ -131,8 +131,10 @@ mapPA_l _ (AClo _fv fl envs) arg@(PNested vsegd _pdata)
  = let  argFlat         = concatPR arg
         c               = lengthPR argFlat
 
-        vseglens        = U.takeLengthsOfVSegd vsegd
-        envsReplicated  = replicatesPR vseglens envs
+        -- TODO: rename this as unsafeDemoteToSegdOfVSegd.. it might overflow
+        segd            = U.demoteToSegdOfVSegd vsegd
+
+        envsReplicated  = replicatesPR segd envs
         arrResult       = fl c envsReplicated argFlat
 
   in    unconcatPR arg arrResult
index e63725b..9a2915e 100644 (file)
@@ -102,7 +102,7 @@ replicatePA n x
 {-# INLINE_PA replicatesPA #-}
 replicatesPA :: PA a => U.Array Int -> PArray a -> PArray a
 replicatesPA repCounts (PArray _ darr)
-        = PArray (U.sum repCounts) (replicatesPR repCounts darr)
+        = PArray (U.sum repCounts) (replicatesPR (U.lengthsToSegd repCounts) darr)
 
 
 -- | Convert a `Vector` to a `PArray`
index aa3b510..9558bc8 100644 (file)
@@ -93,9 +93,13 @@ class PR a where
                 -> a            -- ^ element to replicate.
                 -> PData a
 
-  -- | Segmented replicate.
-  --   O(sum lengths). 
-  replicatesPR       :: U.Array Int -> PData a -> PData a
+  -- | O(sum lengths). Segmented replicate.
+  --   NOTE: This takes a whole Segd instead of just the lengths, because we can
+  --   do it more efficiently if we know there are no zero lengths.
+  --   TODO: the Segd should actually keep track of whether there are zero lengths.
+  replicatesPR  :: U.Segd               -- ^ segment descriptor defining the lengths of the segments.
+                -> PData a              -- ^ data elements to replicate
+                -> PData a
 
   -- | Lookup a single element from the source array.
   --   O(1). 
index 8e2e435..d925a1f 100644 (file)
@@ -56,8 +56,8 @@ instance PR Double where
         = PDouble (U.replicate len x)
 
   {-# INLINE_PDATA replicatesPR #-}
-  replicatesPR lens (PDouble arr)
-        = PDouble (U.replicate_s (U.lengthsToSegd lens) arr)
+  replicatesPR segd (PDouble arr)
+        = PDouble (U.replicate_s segd arr)
                 
   {-# INLINE_PDATA indexPR #-}
   indexPR (PDouble arr) ix
index 484c699..73b8283 100644 (file)
@@ -51,8 +51,8 @@ instance PR Int where
         = PInt (U.replicate len x)
 
   {-# INLINE_PDATA replicatesPR #-}
-  replicatesPR lens (PInt arr)
-        = PInt (U.replicate_s (U.lengthsToSegd lens) arr)
+  replicatesPR segd (PInt arr)
+        = PInt (U.replicate_s segd arr)
                 
   {-# INLINE_PDATA indexPR #-}
   indexPR (PInt arr) ix
index 70aa3a5..e17be80 100644 (file)
@@ -217,7 +217,6 @@ instance PR a => PR (PArray a) where
 
   -- When replicating an array we use the source as the single physical
   -- segment, then point all the virtual segments to it.
-  {-# NOINLINE replicatePR #-}
   replicatePR c (PArray n darr)
    = {-# SCC "replicatePR" #-}
      checkNotEmpty "replicatePR[PArray]" c
@@ -228,16 +227,24 @@ instance PR a => PR (PArray a) where
          uvsegd  = U.mkVSegd (U.replicate c 0) ussegd
 
      in  PNested uvsegd (V.singleton darr)
+  {-# NOINLINE replicatePR #-}
+  --  NOINLINE because it's a cheap segment descriptor operation, 
+  --  and doesn't need to fuse with anything.
                 
 
   -- For segmented replicates, we just replicate the vsegids field.
+  --
   -- TODO: Does replicate_s really need the whole segd,
   --       or could we get away without creating the indices field?
+  --
+  -- TODO: If we know the lens does not contain zeros, then we don't need
+  --       to cull down the psegs.
+  --
   {-# INLINE_PDATA replicatesPR #-}
-  replicatesPR lens (PNested uvsegd pdata)
-   = let segd   = U.lengthsToSegd lens
-     in   PNested (U.updateVSegsOfVSegd (\vsegids -> U.replicate_s segd vsegids) uvsegd)
-                  pdata
+  replicatesPR segd (PNested uvsegd pdata)
+   = PNested (U.updateVSegsOfVSegd      -- TODO use updateReachable if there are no zero len segments.
+                (\vsegids -> U.replicate_s segd vsegids) uvsegd)
+             pdata  
 
 
   -- To index into a nested array, first determine what segment the index
@@ -482,7 +489,6 @@ instance PR a => PR (PArray a) where
 
 
 -------------------------------------------------------------------------------
-
 -- | O(len result). Concatenate a nested array.
 --
 --   This physically performs a 'gather' operation, whereby array data is copied
@@ -525,6 +531,8 @@ concatPR' (PNested vsegd pdatas)
           in   extractsPR pdatas ussegd
 
 {-# NOINLINE concatPR  #-}
+--  TODO: we'll need to inline this when we take the second branch, 
+--  to get specialisation for extractsPR.
 
 
 -- | Build a nested array given a single flat data vector, 
@@ -536,20 +544,24 @@ concatPR' (PNested vsegd pdatas)
 --   segmentation of the template array.
 --
 unconcatPR :: PR a => PData (PArray a) -> PData b -> PData (PArray b)
-unconcatPR arr1 arr2 = {-# SCC "unconcatPR" #-} unconcatPR' arr1 arr2
-unconcatPR' (PNested vsegd pdatas) arr
- = let  
-        -- Get the lengths of all the vsegs individually.
-        !vseglens       = U.takeLengthsOfVSegd vsegd
-
-        -- Rebuild the segd based on these lengths, 
-        -- the resulting segd contains no sharing.
-        !vsegd'         = U.promoteSegdToVSegd
-                        $ U.lengthsToSegd vseglens
+unconcatPR (PNested vsegd pdatas) arr
+ = {-# SCC "unconcatPR" #-}
+   let  
+        -- Demote the vsegd to a manifest vsegd so it contains all the segment
+        -- lengths individually without going through the vsegids.
+        !segd           = U.demoteToSegdOfVSegd vsegd
+
+        -- Rebuild the vsegd based on the manifest vsegd. 
+        -- The vsegids will be just [0..len-1], but this field is constructed
+        -- lazilly and consumers aren't required to demand it.
+        !vsegd'         = U.promoteSegdToVSegd segd
 
    in   PNested vsegd' (V.singleton arr)
 
 {-# NOINLINE unconcatPR #-}
+--  NOINLINE because it won't fuse with anything.
+--  The operations is also entierly on the segment descriptor, so we don't 
+--  need to inline it to specialise it for the element type.
 
 
 -- | Lifted concat.
@@ -568,6 +580,7 @@ concatlPR arr
 
 {-# INLINE concatlPR #-}
 
+
 -- | Lifted append.
 --   Both arrays must contain the same number of elements.
 appendlPR :: PR a => PData (PArray a) -> PData (PArray a) -> PData (PArray a)
index 8f471b6..6bb7aa3 100644 (file)
@@ -51,8 +51,8 @@ instance PR () where
         = PUnit n
 
   {-# INLINE_PDATA replicatesPR #-}
-  replicatesPR lens _
-        = PUnit (U.sum lens)
+  replicatesPR segd _
+        = PUnit (U.elementsSegd segd)
         
   {-# INLINE_PDATA indexPR #-}
   indexPR _ _