Main.hs

module Main where

import Control.Monad
import Criterion.Main
import Data.Char
import Data.Word
import HaskellWorks.Data.Vector.AsVector64s

import qualified Data.ByteString.Lazy                as LBS
import qualified Data.List                           as L
import qualified Data.Succinct.Simd.Capabilities     as CAP
import qualified Data.Succinct.Simd.Comparison.Avx2  as AVX2
import qualified Data.Succinct.Simd.Comparison.Stock as STOCK
import qualified Data.Succinct.Simd.Logical.Avx2     as AVX2
import qualified Data.Succinct.Simd.Logical.Stock    as STOCK
import qualified Data.Vector.Storable                as DVS
import qualified System.Directory                    as IO

runCmpEqWord8sAvx2 :: LBS.ByteString -> IO [DVS.Vector Word64]
runCmpEqWord8sAvx2 bs = do
  let vs = asVector64s 64 bs
  return $ if CAP.avx2Enabled
    then AVX2.cmpEqWord8s (fromIntegral (ord '8')) <$> vs
    else []

runCmpEqWord8sStock :: LBS.ByteString -> IO [DVS.Vector Word64]
runCmpEqWord8sStock bs = do
  let vs = asVector64s 64 bs
  return $ STOCK.cmpEqWord8s (fromIntegral (ord '8')) <$> vs

runAndBitsAvx2 :: LBS.ByteString -> IO [DVS.Vector Word64]
runAndBitsAvx2 bs = do
  let vs = asVector64s 64 bs
  return $ (\v -> AVX2.andBits v v) <$> vs

runAndBitsStock :: LBS.ByteString -> IO [DVS.Vector Word64]
runAndBitsStock bs = do
  let vs = asVector64s 64 bs
  return $ (\v -> STOCK.andBits v v) <$> vs

benchcmpEqWord8s :: IO [Benchmark]
benchcmpEqWord8s = do
  entries <- IO.listDirectory "components/succinct-simd/data/bench"
  let files = ("components/succinct-simd/data/bench/" ++) <$> (".csv" `L.isSuffixOf`) `filter` entries
  benchmarks <- forM files $ \file -> return $
    [ env (LBS.readFile file) $ \bs -> bgroup file
      [ bench ("hw-simd/cmpEqWord8s/avx2/"  <> file) (nfIO (runCmpEqWord8sAvx2  bs))
      , bench ("hw-simd/cmpEqWord8s/stock/" <> file) (nfIO (runCmpEqWord8sStock bs))
      ]
    ]
  return (join benchmarks)

benchAndBits :: IO [Benchmark]
benchAndBits = do
  entries <- IO.listDirectory "components/succinct-simd/data/bench"
  let files = ("components/succinct-simd/data/bench/" ++) <$> (".csv" `L.isSuffixOf`) `filter` entries
  benchmarks <- forM files $ \file -> return $
    [ env (LBS.readFile file) $ \bs -> bgroup file
      [ bench ("hw-simd/andBits/avx2/"  <> file) (nfIO (runAndBitsAvx2  bs))
      , bench ("hw-simd/andBits/stock/" <> file) (nfIO (runAndBitsStock bs))
      ]
    ]
  return (join benchmarks)

main :: IO ()
main = do
  benchmarks <- (mconcat <$>) $ sequence $ mempty
    <> [benchcmpEqWord8s]
    <> [benchAndBits]
  defaultMain benchmarks