[commit: vector] simd: Add a hack to get rid of duplicate loop counters in mzipwith. (17c4b0d)

[Geoffrey Mainland]

Repository : http://darcs.haskell.org/ghc.git/

On branch  : simd

http://hackage.haskell.org/trac/ghc/changeset/17c4b0d9932ce234e49ca9314511e9eb6709da43

>---------------------------------------------------------------

commit 17c4b0d9932ce234e49ca9314511e9eb6709da43
Author: Geoffrey Mainland <mainland at apeiron.net>
Date:   Thu Nov 1 13:08:15 2012 +0000

    Add a hack to get rid of duplicate loop counters in mzipwith.

>---------------------------------------------------------------

 Data/Vector/Fusion/Bundle/Monadic.hs |   45 ++++++++++++++++++++++++++++++++++
 Data/Vector/Generic.hs               |    7 ++++++
 2 files changed, 52 insertions(+)

diff --git a/Data/Vector/Fusion/Bundle/Monadic.hs b/Data/Vector/Fusion/Bundle/Monadic.hs
index cf10361..8954d04 100644
--- a/Data/Vector/Fusion/Bundle/Monadic.hs
+++ b/Data/Vector/Fusion/Bundle/Monadic.hs
@@ -83,6 +83,7 @@ module Data.Vector.Fusion.Bundle.Monadic (
 
   -- * Multi Zipping
   mzipWithM_, mzipWithM, mzipWith,
+  mzipWithHackM, mzipWithHack,
 
   -- * Multi Folding
   mfoldl, mfoldlM, mfoldM,
@@ -1294,6 +1295,50 @@ mzipWith :: Monad m
 {-# INLINE mzipWith #-}
 mzipWith p q = mzipWithM (\a b -> return (p a b)) (\a b -> return (q a b))
 
+mzipWithHackM  ::  forall m v a b . (PackedVector v a, Monad m)
+               =>  (a -> a -> m b)
+               ->  (Multi a -> Multi a -> m (Multi b))
+               ->  v a
+               ->  v a
+               ->  Bundle m v b
+{-# INLINE_FUSED mzipWithHackM #-}
+mzipWithHackM p q v1 v2 =
+    v1 `seq` v2 `seq` n `seq`
+      Bundle (Stream step 0)
+             (Left (MultiStream stepm step 0))
+             (Stream vstep 0)
+             Nothing
+             (Exact n)
+  where
+    n = min (basicLength v1) (basicLength v2)
+    k = n - n `rem` m
+    m = multiplicity (undefined :: Multi a)
+
+    {-# INLINE step #-}
+    step i | i >= n    = return Done
+           | otherwise = case basicUnsafeIndexM v1 i of
+                           Box x -> case basicUnsafeIndexM v2 i of
+                                      Box y -> liftM (`Yield` (i+1)) (p x y)
+
+    {-# INLINE stepm #-}
+    stepm i | i >= k    = return Done
+            | otherwise = case basicUnsafeIndexAsMultiM v1 i of
+                            Box x -> case basicUnsafeIndexAsMultiM v2 i of
+                                       Box y -> liftM (`Yield` (i+m)) (q x y)
+    
+    {-# INLINE vstep #-}
+    vstep i = do r <- step i
+                 return $ fmap (\x -> Chunk 1 (\v -> M.basicUnsafeWrite v 0 x)) r
+
+mzipWithHack :: (PackedVector v a, Monad m)
+             => (a -> a -> b)
+             -> (Multi a -> Multi a -> Multi b)
+             -> v a
+             -> v a
+             -> Bundle m v b
+{-# INLINE mzipWithHack #-}
+mzipWithHack p q = mzipWithHackM (\a b -> return (p a b)) (\a b -> return (q a b))
+
 -- Multi Folding
 -- -------------
 
diff --git a/Data/Vector/Generic.hs b/Data/Vector/Generic.hs
index 7135906..4d19756 100644
--- a/Data/Vector/Generic.hs
+++ b/Data/Vector/Generic.hs
@@ -2160,6 +2160,13 @@ mzipWithM_ :: (Monad m, PackedVector v a, PackedVector v b, PackedVector v c)
 {-# INLINE mzipWithM_ #-}
 mzipWithM_ p q as bs = Bundle.mzipWithM_ p q (multistream as) (multistream bs)
 
+{-# RULES
+
+"mzipWithM/multistream [Vector]" [2] forall p q v1 v2 .
+  MBundle.mzipWithM p q (multistream v1) (multistream v2) = MBundle.mzipWithHackM p q v1 v2
+
+  #-}
+
 mfoldl :: PackedVector v b => (a -> b -> a) -> (a -> Multi b -> a) -> a -> v b -> a
 {-# INLINE mfoldl #-}
 mfoldl p q z = Bundle.mfoldl p q z . multistream