pure-borrow-0.1.0.0: bench/qsort.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE QualifiedDo #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}

module Main (main) where

import Control.Concurrent (getNumCapabilities)
import Control.Concurrent.DivideConquer.Linear (
  naiveDivideAndConquer,
  qsort,
  qsortDC,
  qsortDC',
 )
import Control.Functor.Linear qualified as Control
import Control.Monad.Borrow.Pure.BO
import Control.Syntax.DataFlow qualified as DataFlow
import Data.Proxy (Proxy (..))
import Data.Vector.Algorithms.Intro qualified as AI
import Data.Vector.Generic.Mutable.Linear.Borrow.Unrestricted qualified as VL
import Data.Vector.Unboxed qualified as V
import Prelude.Linear (dup, unur)
import Prelude.Linear qualified as PL
import System.Random.Stateful
import Test.Tasty (askOption, defaultMainWithIngredients)
import Test.Tasty.Bench
import Test.Tasty.Ingredients.Basic (includingOptions)
import Test.Tasty.Options
import Text.Read (readMaybe)
import Prelude as P

data Mode = Parallel Word | NaiveDC Int | Worksteal Int | Sequential | IntroSort
  deriving (Show, Eq, Ord)

data BenchOpts = BenchOpts {numThreads :: !Int, sampleSize :: !Int}
  deriving (Show, Eq, Ord)

qsortWith :: Mode -> V.Vector Int -> V.Vector Int
qsortWith IntroSort v = V.modify AI.sort v
qsortWith (Parallel budget) v =
  unur PL.$ linearly \lin ->
    DataFlow.do
      (lin, l2) <- dup lin
      runBO lin Control.do
        (v, lend) <- borrowM (VL.fromVector v l2)
        qsort budget v
        Control.pure PL.$ VL.toVector Control.<$> reclaim' lend
qsortWith Sequential v =
  unur PL.$ linearly \lin ->
    DataFlow.do
      (lin, l2) <- dup lin
      runBO lin Control.do
        (v, lend) <- borrowM (VL.fromVector v l2)
        qsort 0 v
        pureAfter (VL.toVector PL.$ reclaim lend)
qsortWith (NaiveDC p) v =
  unur PL.$ linearly \lin ->
    DataFlow.do
      (lin, l2) <- dup lin
      runBO lin Control.do
        (v, lend) <- borrowM (VL.fromVector v l2)
        Control.void PL.$ naiveDivideAndConquer (qsortDC' p) v
        pureAfter (VL.toVector PL.$ reclaim lend)
qsortWith (Worksteal p) v =
  unur PL.$ linearly \lin ->
    DataFlow.do
      (lin, l2) <- dup lin
      runBO lin Control.do
        (v, lend) <- borrowM (VL.fromVector v l2)
        Control.void PL.$ qsortDC (mkStdGen 42) p 128 v
        pureAfter (VL.toVector PL.$ reclaim lend)

data SampleSize = SampleSize Int
  deriving (Show, Eq, Ord)

instance IsOption SampleSize where
  defaultValue = SampleSize 32
  parseValue s =
    case readMaybe s of
      Just n | n > 0, kMAX_SIZE `rem` n == 0 -> Just (SampleSize n)
      _ -> Nothing
  optionName = return "size"
  optionHelp =
    return "Number of linear size steps (positive divisor of 32768)"

main :: IO ()
main = do
  numThreads <- getNumCapabilities
  let customOpts = [Option (Proxy :: Proxy SampleSize)]
      ingredients = includingOptions customOpts : benchIngredients
  defaultMainWithIngredients ingredients $ askOption \(SampleSize sampleSize) ->
    bgroup "All" $ benches BenchOpts {..}

benches :: BenchOpts -> [Benchmark]
benches BenchOpts {..} =
  [ bgroup
      "qsort"
      [ env
          ( pure $ runStateGen_ (mkStdGen 42) \g -> do
              V.replicateM size (uniformM g)
          )
          \vec ->
            bgroup
              (show size)
              ( [ bench "intro" $ nf (qsortWith IntroSort) vec
                , bench "sequential" $ nf (qsortWith Sequential) vec
                ]
                  <> [ bench ("parallel (budget = " <> show n <> ")") $
                         nf (qsortWith $ Parallel n) vec
                     | n <- [4, 8, 16, 32]
                     ]
                  <> [ bench ("parallel-dc (thresh = 128)") $
                         nf (qsortWith $ NaiveDC 128) vec
                     ]
                  <> [ bench ("worksteal (workers = " <> show n <> ")") $
                         nf (qsortWith $ Worksteal n) vec
                     | n <- [2, 4 .. numThreads]
                     ]
              )
      | i <- [0 .. sampleSize]
      , let size = i * kMAX_SIZE `quot` sampleSize
      ]
  ]

kMAX_SIZE :: Int
kMAX_SIZE = 32 * 1024