Matrix.hs

{-# LANGUAGE DataKinds #-}

module Bench.Matrix (benches) where

import AtCoder.Extra.Math qualified as ACEM
import AtCoder.Extra.Semigroup.Matrix qualified as ACMAT
import AtCoder.ModInt qualified as M
import BenchLib.Matrix qualified as Mat
import Control.Monad.State.Class (MonadState, state)
import Control.Monad.Trans.State.Strict (evalState, runState)
import Criterion
import Data.Semigroup (Semigroup (..))
import Data.Vector qualified as V
import Data.Vector.Unboxed qualified as VU
import System.Random

testSize :: Int
testSize = 10000

m :: Int
m = 998244353

randomMatrix :: (MonadState StdGen m) => m (Mat.Matrix Int, Int)
randomMatrix = do
  h <- state $ uniformR (1, 16)
  nextMatrix h

nextMatrix :: (MonadState StdGen m) => Int -> m (Mat.Matrix Int, Int)
nextMatrix h = do
  w <- state $ uniformR (1, 16)
  vec <- VU.replicateM (h * w) $ state (uniformR (0, m - 1))
  pure (Mat.new h w vec, w)

randomSquareACLMatrix :: (MonadState StdGen m) => Int -> m (ACMAT.Matrix Int)
randomSquareACLMatrix n = do
  vec <- VU.replicateM (n * n) $ state (uniformR (0, m - 1))
  pure $ ACMAT.new n n vec

benches :: Benchmark
benches =
  bgroup
    "Matrix"
    [ bench "mul1" $ whnf (V.foldl1' Mat.mul1) randomMatrixInput,
      bench "mul2" $ whnf (V.foldl1' Mat.mul2) randomMatrixInput,
      bench "mul3_1" $ whnf (V.foldl1' Mat.mul3_1) randomMatrixInput,
      bench "mul3_2" $ whnf (V.foldl1' Mat.mul3_2) randomMatrixInput,
      bench "mul3_3" $ whnf (V.foldl1' Mat.mul3_3) randomMatrixInput,
      bench "mul4_1" $ whnf (V.foldl1' Mat.mul4_1) randomMatrixInput,
      bench "mul4_2" $ whnf (V.foldl1' Mat.mul4_2) randomMatrixInput,
      bench "mul4_3" $ whnf (V.foldl1' Mat.mul4_3) randomMatrixInput,
      bench "mulMod1" $ whnf (V.foldl1' (Mat.mulMod1 m)) randomMatrixInput,
      bench "mulMod2" $ whnf (V.foldl1' (Mat.mulMod2 m)) randomMatrixInput,
      bench "mulMod3" $ whnf (V.foldl1' (Mat.mulMod3 m)) randomMatrixInput,
      bench "mulMod4" $ whnf (V.foldl1' (Mat.mulMod4 m)) randomMatrixInput,
      bench "mulMod5" $ whnf (V.foldl1' (Mat.mulMod5 m)) randomMatrixInput,
      bench "mulMod6" $ whnf (V.foldl1' (Mat.mulMod6 m)) randomMatrixInput,
      bench "mulMod7" $ whnf (V.foldl1' (Mat.mulMod7 m)) randomMatrixInput,
      bench "mulMint1" $ whnf (V.foldl1' Mat.mulMint1) randomMintMatrixInput,
      bench "mulMint2" $ whnf (V.foldl1' Mat.mulMint2) randomMintMatrixInput,
      bench "mulMint3" $ whnf (V.foldl1' Mat.mulMint3) randomMintMatrixInput,
      bench "mulMint4" $ whnf (V.foldl1' Mat.mulMint4) randomMintMatrixInput,
      -- mul (ACL)
      bench "mul_ACL" $ whnf (V.foldl1' ACMAT.mul) randomMatrixInputACL,
      bench "mulMod_ACL" $ whnf (V.foldl1' (ACMAT.mulMod m)) randomMatrixInputACL,
      bench "mulMint_ACL" $ whnf (V.foldl1' ACMAT.mulMint) randomMintMatrixInputACL,
      -- pow mod (ACL only)
      bench "powMod_ACL" $ whnf (VU.foldl' (flip (ACMAT.powMod m)) squareMat) randomVec,
      bench "powMintACL_stimes" $ whnf (VU.foldl' (flip stimes) squareMatMint) randomVec,
      bench "powMintACL_stimes'" $ whnf (VU.foldl' (flip ACEM.stimes') squareMatMint) randomVec,
      bench "powMintACL_powMint" $ whnf (VU.foldl' (flip ACMAT.powMint) squareMatMint) randomVec
    ]
  where
    -- Bench matrix
    randomMatrixInput :: V.Vector (Mat.Matrix Int)
    !randomMatrixInput = evalState (V.unfoldrExactNM testSize nextMatrix (Mat.wM mat0)) gen0
      where
        ((!mat0, !_), !gen0) = runState randomMatrix $ mkStdGen 123456789

    randomMintMatrixInput :: V.Vector (Mat.Matrix (M.ModInt 998244353))
    !randomMintMatrixInput = V.map (Mat.map M.new) randomMatrixInput

    -- ACL matrix
    randomMatrixInputACL :: V.Vector (ACMAT.Matrix Int)
    !randomMatrixInputACL = V.map (\(Mat.Matrix h w vec) -> ACMAT.new h w vec) randomMatrixInput

    randomMintMatrixInputACL :: V.Vector (ACMAT.Matrix (M.ModInt 998244353))
    !randomMintMatrixInputACL = V.map (\(Mat.Matrix h w vec) -> ACMAT.new h w vec) randomMintMatrixInput

    squareMat :: ACMAT.Matrix Int
    !squareMat = evalState (randomSquareACLMatrix 17) $ mkStdGen 123456789

    squareMatMint :: ACMAT.Matrix (M.ModInt 998244353)
    !squareMatMint = ACMAT.map M.new squareMat

    -- non-zero random vector
    randomVec :: VU.Vector Int
    !randomVec =
      VU.map ((+ 1) . fromIntegral) $
        VU.unfoldrExactN 100 (genWord64R (998244353 - 2)) (mkStdGen 123456789)