Add more iter benchmarks

lehins · Aug 7, 2022 · f855e94 · f855e94
1 parent f96d02a
commit f855e94
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 26 deletions.
diff --git a/massiv-bench/bench/Iter.hs b/massiv-bench/bench/Iter.hs
@@ -3,10 +3,10 @@
 {-# LANGUAGE FlexibleContexts #-}
 module Main where
 
-import Criterion.Main
-import Data.Massiv.Array as A
 import Control.Monad.ST
 import Control.Scheduler
+import Criterion.Main
+import Data.Massiv.Array as A
 import Data.Typeable
 
 baseline :: Sz1 -> IO ()
@@ -67,20 +67,22 @@ iterFullBench :: Index ix => Sz ix -> Benchmark
 iterFullBench !sz =
   bgroup
     (show (typeOf sz))
-    [ bench "iterA_" $ whnfIO $
-      iterA_ zeroIndex (unSz sz) oneIndex (<) (\ix -> ix `seq` pure ())
-    , bench "iterFullA_ (RowMajor)" $ whnfIO $
+    [ bench "iterFullA_ (RowMajor)" $ whnfIO $
       iterFullA_ defRowMajor zeroIndex sz (\ix -> ix `seq` pure ())
     , bench "iterTargetFullST_ (RowMajor)" $ whnfIO $
-      stToIO $ iterTargetFullST_ defRowMajor trivialScheduler_ 0 sz seqAction
+      stToIO $ iterTargetFullST_ defRowMajor trivialScheduler_ 0 sz seq2Action
+    , bench "iterFullM (RowMajor)" $ whnfIO $
+      stToIO $ iterFullM defRowMajor zeroIndex sz () seq2Action
     , bench "iterFullST (RowMajor)" $ whnfIO $
-      stToIO $ iterFullST defRowMajor trivialScheduler_ zeroIndex sz () noopSplit seqAction
+      stToIO $ iterFullST defRowMajor trivialScheduler_ zeroIndex sz () noopSplit seq2Action
     , bench "iterFullA_ (RowMajorLinear)" $ whnfIO $
       iterFullA_ defRowMajorLinear zeroIndex sz (\ix -> ix `seq` pure ())
     , bench "iterTargetFullST_ (RowMajorLinear)" $ whnfIO $
-      stToIO $ iterTargetFullST_ defRowMajorLinear trivialScheduler_ 0 sz seqAction
+      stToIO $ iterTargetFullST_ defRowMajorLinear trivialScheduler_ 0 sz seq2Action
+    , bench "iterFullM (RowMajorLinear)" $ whnfIO $
+      stToIO $ iterFullM defRowMajorLinear zeroIndex sz () seq2Action
     , bench "iterFullST (RowMajorLinear)" $ whnfIO $
-      stToIO $ iterFullST defRowMajorLinear trivialScheduler_ zeroIndex sz () noopSplit seqAction
+      stToIO $ iterFullST defRowMajorLinear trivialScheduler_ zeroIndex sz () noopSplit seq2Action
     ]
 
 noopSplit :: () -> ST s ((), ())
@@ -90,18 +92,13 @@ iterStrideBench :: (Num ix, Index ix) => Stride ix -> Sz ix -> Benchmark
 iterStrideBench !stride !sz =
   bgroup
     (show (typeOf sz))
-    [ bench "iterM" $
+    [ bench "iterTargetFullWithStrideST_ (RowMajor)" $
       whnfIO $
-      () <$ iterM zeroIndex (unSz sz * s) s (<) (0 :: Int) (\ix i -> (i + 1) <$ seqAction i ix)
-    , bench "iterTargetFullWithStrideST_ (RowMajor)" $
-      whnfIO $
-      stToIO $ iterTargetFullWithStrideST_ defRowMajor trivialScheduler_ 0 sz stride seqAction
+      stToIO $ iterTargetFullWithStrideST_ defRowMajor trivialScheduler_ 0 sz stride seq2Action
     , bench "iterTargetFullWithStrideST_ (RowMajorLinear)" $
       whnfIO $
-      stToIO $ iterTargetFullWithStrideST_ defRowMajorLinear trivialScheduler_ 0 sz stride seqAction
+      stToIO $ iterTargetFullWithStrideST_ defRowMajorLinear trivialScheduler_ 0 sz stride seq2Action
     ]
-  where
-    s = unStride stride
 
 
 iterFullBenchPar :: Index ix => Sz ix -> Benchmark
@@ -110,16 +107,16 @@ iterFullBenchPar sz =
     (show (typeOf sz))
     [ bench "iterTargetFullST_ (RowMajor)" $ whnfIO $
       withMassivScheduler_ Par $ \ scheduler ->
-        stToIO $ iterTargetFullST_ defRowMajor scheduler 0 sz seqAction
+        stToIO $ iterTargetFullST_ defRowMajor scheduler 0 sz seq2Action
     , bench "iterFullST (RowMajor)" $ whnfIO $
       withMassivScheduler_ Par $ \ scheduler ->
-        stToIO $ iterFullST defRowMajor scheduler zeroIndex sz () noopSplit seqAction
+        stToIO $ iterFullST defRowMajor scheduler zeroIndex sz () noopSplit seq2Action
     , bench "iterTargetFullST_ (RowMajorLinear)" $ whnfIO $
       withMassivScheduler_ Par $ \ scheduler ->
-        stToIO $ iterTargetFullST_ defRowMajorLinear scheduler 0 sz seqAction
+        stToIO $ iterTargetFullST_ defRowMajorLinear scheduler 0 sz seq2Action
     , bench "iterFullST (RowMajorLinear)" $ whnfIO $
       withMassivScheduler_ Par $ \ scheduler ->
-        stToIO $ iterFullST defRowMajorLinear scheduler zeroIndex sz () noopSplit seqAction
+        stToIO $ iterFullST defRowMajorLinear scheduler zeroIndex sz () noopSplit seq2Action
     ]
 
 iterStrideBenchPar :: (Num ix, Index ix) => Stride ix -> Sz ix -> Benchmark
@@ -128,15 +125,26 @@ iterStrideBenchPar stride sz =
     (show (typeOf sz))
     [ bench "iterTargetFullWithStrideST_ (RowMajor)" $ whnfIO $
       withMassivScheduler_ Par $ \ scheduler ->
-        stToIO $ iterTargetFullWithStrideST_ defRowMajor scheduler 0 sz stride seqAction
+        stToIO $ iterTargetFullWithStrideST_ defRowMajor scheduler 0 sz stride seq2Action
+    , bench "iterTargetFullWithStrideST (RowMajor)" $ whnfIO $
+      withMassivScheduler_ Par $ \ scheduler ->
+        stToIO $ iterTargetFullWithStrideST defRowMajor scheduler 0 sz stride () noopSplit seq3Action
     , bench "iterTargetFullWithStrideST_ (RowMajorLinear)" $ whnfIO $
       withMassivScheduler_ Par $ \ scheduler ->
-        stToIO $ iterTargetFullWithStrideST_ defRowMajorLinear scheduler 0 sz stride seqAction
+        stToIO $ iterTargetFullWithStrideST_ defRowMajorLinear scheduler 0 sz stride seq2Action
+    , bench "iterTargetFullWithStrideST (RowMajorLinear)" $ whnfIO $
+      withMassivScheduler_ Par $ \ scheduler ->
+        stToIO $
+        iterTargetFullWithStrideST defRowMajorLinear scheduler 0 sz stride () noopSplit seq3Action
     ]
 
-seqAction :: Monad m => a -> b -> m ()
-seqAction a b = a `seq` b `seq` pure ()
-{-# INLINE seqAction #-}
+seq2Action :: Monad m => a -> b -> m ()
+seq2Action a b = a `seq` b `seq` pure ()
+{-# INLINE seq2Action #-}
+
+seq3Action :: Monad m => a -> b -> c -> m ()
+seq3Action a b c = a `seq` b `seq` c `seq` pure ()
+{-# INLINE seq3Action #-}
 
 
 
diff --git a/massiv/src/Data/Massiv/Core/Index/Ix.hs b/massiv/src/Data/Massiv/Core/Index/Ix.hs
@@ -423,6 +423,13 @@ instance {-# OVERLAPPING #-} Index (IxN 3) where
           let (i3 :> i2 :. i1) = fromLinearIndex sz i
           in action (iShift + i) ((b3 + s3 * i3) :> (b2 + s2 * i2) :. (b1 + s1 * i1))
   {-# INLINE iterTargetRowMajorAccST_ #-}
+  iterTargetRowMajorAccST iAcc fact scheduler iStart sz (b3 :> b2 :. b1) (s3 :> s2 :. s1) acc splitAcc action =
+    let n = totalElem sz
+        iShift = iStart + iAcc * n
+     in iterLinearAccST fact scheduler 0 1 n acc splitAcc $ \ !i ->
+          let (i3 :> i2 :. i1) = fromLinearIndex sz i
+          in action (iShift + i) ((b3 + s3 * i3) :> (b2 + s2 * i2) :. (b1 + s1 * i1))
+  {-# INLINE iterTargetRowMajorAccST #-}
 
 -- | Constraint synonym that encapsulates all constraints needed for dimension 4 and higher.
 --