Bench.hs

-- | Facilities for handling Futhark benchmark results.  A Futhark
-- benchmark program is just like a Futhark test program.
module Futhark.Bench
  ( RunResult (..),
    DataResult (..),
    BenchResult (..),
    Result (..),
    encodeBenchResults,
    decodeBenchResults,
    binaryName,
    benchmarkDataset,
    RunOptions (..),
    prepareBenchmarkProgram,
    CompileOptions (..),
    module Futhark.Profile,
  )
where

import Control.Applicative
import Control.Monad
import Control.Monad.Except (ExceptT, MonadError (..), liftEither, runExceptT)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Aeson qualified as JSON
import Data.Aeson.Key qualified as JSON
import Data.Aeson.KeyMap qualified as JSON
import Data.ByteString.Char8 qualified as SBS
import Data.ByteString.Lazy.Char8 qualified as LBS
import Data.DList qualified as DL
import Data.Map qualified as M
import Data.Maybe
import Data.Text qualified as T
import Data.Time.Clock
import Data.Vector.Unboxed qualified as U
import Futhark.Profile
import Futhark.Server
import Futhark.Test
import Futhark.Util (showText)
import Statistics.Autocorrelation (autocorrelation)
import Statistics.Sample (fastStdDev, mean)
import System.Exit
import System.FilePath
import System.Process.ByteString (readProcessWithExitCode)
import System.Timeout (timeout)

-- | The runtime of a single succesful run.
newtype RunResult = RunResult {runMicroseconds :: Int}
  deriving (Eq, Show)

-- | The measurements resulting from various successful runs of a
-- benchmark (same dataset).
data Result = Result
  { -- | Runtime of every run.
    runResults :: [RunResult],
    -- | Memory usage.
    memoryMap :: M.Map T.Text Int,
    -- | The error output produced during execution.  Often 'Nothing'
    -- for space reasons, and otherwise only the output from the last
    -- run.
    stdErr :: Maybe T.Text,
    -- | Profiling report.  This will have been measured based on the
    -- last run.
    report :: Maybe ProfilingReport
  }
  deriving (Eq, Show)

-- | The results for a single named dataset is either an error message, or
-- a result.
data DataResult = DataResult T.Text (Either T.Text Result)
  deriving (Eq, Show)

-- | The results for all datasets for some benchmark program.
data BenchResult = BenchResult
  { benchResultProg :: FilePath,
    benchResultResults :: [DataResult]
  }
  deriving (Eq, Show)

newtype DataResults = DataResults {unDataResults :: [DataResult]}

newtype BenchResults = BenchResults {unBenchResults :: [BenchResult]}

instance JSON.ToJSON Result where
  toJSON (Result runres memmap err profiling) =
    JSON.toJSON (runres, memmap, err, profiling)

instance JSON.FromJSON Result where
  parseJSON = fmap f . JSON.parseJSON
    where
      f (runres, memmap, err, profiling) = Result runres memmap err profiling

instance JSON.ToJSON RunResult where
  toJSON = JSON.toJSON . runMicroseconds

instance JSON.FromJSON RunResult where
  parseJSON = fmap RunResult . JSON.parseJSON

instance JSON.ToJSON DataResults where
  toJSON (DataResults rs) =
    JSON.object $ map dataResultJSON rs
  toEncoding (DataResults rs) =
    JSON.pairs $ mconcat $ map (uncurry (JSON..=) . dataResultJSON) rs

instance JSON.FromJSON DataResults where
  parseJSON = JSON.withObject "datasets" $ \o ->
    DataResults <$> mapM datasetResult (JSON.toList o)
    where
      datasetResult (k, v) =
        DataResult (JSON.toText k)
          <$> ((Right <$> success v) <|> (Left <$> JSON.parseJSON v))
      success = JSON.withObject "result" $ \o ->
        Result
          <$> o JSON..: "runtimes"
          <*> o JSON..: "bytes"
          <*> o JSON..:? "stderr"
          <*> o JSON..:? "profiling"

dataResultJSON :: DataResult -> (JSON.Key, JSON.Value)
dataResultJSON (DataResult desc (Left err)) =
  (JSON.fromText desc, JSON.toJSON err)
dataResultJSON (DataResult desc (Right (Result runtimes bytes progerr_opt profiling_opt))) =
  ( JSON.fromText desc,
    JSON.object $
      [ ("runtimes", JSON.toJSON $ map runMicroseconds runtimes),
        ("bytes", JSON.toJSON bytes)
      ]
        <> case progerr_opt of
          Nothing -> []
          Just progerr -> [("stderr", JSON.toJSON progerr)]
        <> case profiling_opt of
          Nothing -> []
          Just profiling -> [("profiling", JSON.toJSON profiling)]
  )

benchResultJSON :: BenchResult -> (JSON.Key, JSON.Value)
benchResultJSON (BenchResult prog r) =
  ( JSON.fromString prog,
    JSON.object [("datasets", JSON.toJSON $ DataResults r)]
  )

instance JSON.ToJSON BenchResults where
  toJSON (BenchResults rs) =
    JSON.object $ map benchResultJSON rs

instance JSON.FromJSON BenchResults where
  parseJSON = JSON.withObject "benchmarks" $ \o ->
    BenchResults <$> mapM onBenchmark (JSON.toList o)
    where
      onBenchmark (k, v) =
        BenchResult (JSON.toString k)
          <$> JSON.withObject "benchmark" onBenchmark' v
      onBenchmark' o =
        fmap unDataResults . JSON.parseJSON =<< o JSON..: "datasets"

-- | Transform benchmark results to a JSON bytestring.
encodeBenchResults :: [BenchResult] -> LBS.ByteString
encodeBenchResults = JSON.encode . BenchResults

-- | Decode benchmark results from a JSON bytestring.
decodeBenchResults :: LBS.ByteString -> Either String [BenchResult]
decodeBenchResults = fmap unBenchResults . JSON.eitherDecode'

--- Running benchmarks

-- | How to run a benchmark.
data RunOptions = RunOptions
  { -- | Applies both to initial and convergence phase.
    runMinRuns :: Int,
    runMinTime :: NominalDiffTime,
    runTimeout :: Int,
    runVerbose :: Int,
    -- | If true, run the convergence phase.
    runConvergencePhase :: Bool,
    -- | Stop convergence once this much time has passed.
    runConvergenceMaxTime :: NominalDiffTime,
    -- | Invoked for every runtime measured during the run.  Can be
    -- used to provide a progress bar.
    runResultAction :: (Int, Maybe Double) -> IO (),
    -- | Perform a final run at the end with profiling information
    -- enabled.
    runProfile :: Bool
  }

-- | A list of @(autocorrelation,rse)@ pairs.  When the
-- autocorrelation is above the first element and the RSE is above the
-- second element, we want more runs.
convergenceCriteria :: [(Double, Double)]
convergenceCriteria =
  [ (0.95, 0.0010),
    (0.75, 0.0015),
    (0.65, 0.0025),
    (0.45, 0.0050),
    (0.00, 0.0100)
  ]

-- Returns the next run count.
nextRunCount :: Int -> Double -> Double -> Int
nextRunCount runs rse acor = if any check convergenceCriteria then div runs 2 else 0
  where
    check (acor_lb, rse_lb) = acor > acor_lb && rse > rse_lb

type BenchM = ExceptT T.Text IO

-- Do the minimum number of runs.
runMinimum ::
  BenchM (RunResult, [T.Text]) ->
  RunOptions ->
  Int ->
  NominalDiffTime ->
  DL.DList (RunResult, [T.Text]) ->
  BenchM (DL.DList (RunResult, [T.Text]))
runMinimum do_run opts runs_done elapsed r = do
  let actions = do
        x <- do_run
        liftIO $ runResultAction opts (runMicroseconds (fst x), Nothing)
        pure x

  -- Figure out how much we have left to do.
  let todo
        | runs_done < runMinRuns opts =
            runMinRuns opts - runs_done
        | otherwise =
            -- Guesstimate how many runs we need to reach the minimum
            -- time.
            let time_per_run = elapsed / fromIntegral runs_done
             in ceiling ((runMinTime opts - elapsed) / time_per_run)

  -- Note that todo might be negative if minimum time has been exceeded.
  if todo <= 0
    then pure r
    else do
      before <- liftIO getCurrentTime
      r' <- DL.fromList <$> replicateM todo actions
      after <- liftIO getCurrentTime
      let elapsed' = elapsed + diffUTCTime after before
      runMinimum do_run opts (runs_done + todo) elapsed' (r <> r')

-- Do more runs until a convergence criterion is reached.
runConvergence ::
  BenchM (RunResult, [T.Text]) ->
  RunOptions ->
  DL.DList (RunResult, [T.Text]) ->
  BenchM (DL.DList (RunResult, [T.Text]))
runConvergence do_run opts initial_r =
  let runtimes = resultRuntimes (DL.toList initial_r)
      (n, _, rse, acor) = runtimesMetrics runtimes
   in -- If the runtimes collected during the runMinimum phase are
      -- unstable enough that we need more in order to converge, we throw
      -- away the runMinimum runtimes.  This is because they often exhibit
      -- behaviour similar to a "warmup" period, and hence function as
      -- outliers that poison the metrics we use to determine convergence.
      -- By throwing them away we converge much faster, and still get the
      -- right result.
      case nextRunCount n rse acor of
        x
          | x > 0,
            runConvergencePhase opts ->
              moreRuns mempty mempty rse (x `max` runMinRuns opts)
          | otherwise ->
              pure initial_r
  where
    resultRuntimes =
      U.fromList . map (fromIntegral . runMicroseconds . fst)

    runtimesMetrics runtimes =
      let n = U.length runtimes
          rse = (fastStdDev runtimes / sqrt (fromIntegral n)) / mean runtimes
          (x, _, _) = autocorrelation runtimes
       in ( n,
            realToFrac (U.sum runtimes) :: NominalDiffTime,
            rse,
            fromMaybe 1 (x U.!? 1)
          )

    sample rse = do
      x <- do_run
      liftIO $ runResultAction opts (runMicroseconds (fst x), Just rse)
      pure x

    moreRuns runtimes r rse x = do
      r' <- replicateM x $ sample rse
      loop (runtimes <> resultRuntimes r') (r <> DL.fromList r')

    loop runtimes r = do
      let (n, total, rse, acor) = runtimesMetrics runtimes
      case nextRunCount n rse acor of
        x
          | x > 0,
            total < runConvergenceMaxTime opts ->
              moreRuns runtimes r rse x
          | otherwise ->
              pure r

-- | Run the benchmark program on the indicated dataset.
benchmarkDataset ::
  Server ->
  RunOptions ->
  FutharkExe ->
  FilePath ->
  T.Text ->
  Values ->
  Maybe Success ->
  FilePath ->
  IO (Either T.Text ([RunResult], T.Text, ProfilingReport))
benchmarkDataset server opts futhark program entry input_spec expected_spec ref_out = runExceptT $ do
  input_types <- cmdEither $ cmdInputs server entry
  output_type <- cmdEither $ cmdOutput server entry
  let out = "out"
      ins = ["in" <> showText i | i <- [0 .. length input_types - 1]]

  cmdMaybe . liftIO $ cmdClear server

  cmdMaybe . liftIO $ cmdPauseProfiling server

  let freeOut = cmdMaybe (cmdFree server [out])
      freeIns = cmdMaybe (cmdFree server ins)
      loadInput = valuesAsVars server (zip ins $ map inputType input_types) futhark dir input_spec
      reloadInput = freeIns >> loadInput

  loadInput

  let runtime l
        | Just l' <- T.stripPrefix "runtime: " l,
          [(x, "")] <- reads $ T.unpack l' =
            Just x
        | otherwise =
            Nothing

      doRun = do
        call_lines <- cmdEither (cmdCall server entry out ins)
        when (any inputConsumed input_types) reloadInput
        case mapMaybe runtime call_lines of
          [call_runtime] -> pure (RunResult call_runtime, call_lines)
          [] -> throwError "Could not find runtime in output."
          ls -> throwError $ "Ambiguous runtimes: " <> showText ls

  maybe_call_logs <- liftIO . timeout (runTimeout opts * 1000000) . runExceptT $ do
    -- First one uncounted warmup run.
    void $ cmdEither $ cmdCall server entry out ins

    ys <- runMinimum (freeOut *> doRun) opts 0 0 mempty

    xs <- runConvergence (freeOut *> doRun) opts ys

    -- Possibly a profiled run at the end.
    profile_log <-
      if not (runProfile opts)
        then pure Nothing
        else do
          cmdMaybe . liftIO $ cmdUnpauseProfiling server
          profile_log <- freeOut *> doRun
          cmdMaybe . liftIO $ cmdPauseProfiling server
          pure $ Just profile_log

    pure (DL.toList xs, profile_log)

  (call_logs, profile_log) <- case maybe_call_logs of
    Nothing ->
      throwError . T.pack $
        "Execution exceeded " ++ show (runTimeout opts) ++ " seconds."
    Just x -> liftEither x

  freeIns

  report <- cmdEither $ cmdReport server

  report' <-
    maybe (throwError "Program produced invalid profiling report.") pure $
      profilingReportFromText (T.unlines report)

  maybe_expected <-
    liftIO $ maybe (pure Nothing) (fmap Just . getExpectedValues) expected_spec

  case maybe_expected of
    Just expected -> do
      vs <- readResults server (out, outputType output_type) <* freeOut
      checkResult program expected vs
    Nothing ->
      freeOut

  pure
    ( map fst call_logs,
      T.unlines $ snd $ fromMaybe (last call_logs) profile_log,
      report'
    )
  where
    getExpectedValues (SuccessValues vs) =
      getValues futhark dir vs
    getExpectedValues SuccessGenerateValues =
      getExpectedValues $ SuccessValues $ InFile ref_out

    dir = takeDirectory program

-- | How to compile a benchmark.
data CompileOptions = CompileOptions
  { compFuthark :: String,
    compBackend :: String,
    compOptions :: [String]
  }

progNotFound :: String -> String
progNotFound s = s ++ ": command not found"

-- | Compile and produce reference datasets.
prepareBenchmarkProgram ::
  (MonadIO m) =>
  Maybe Int ->
  CompileOptions ->
  FilePath ->
  [InputOutputs] ->
  m (Either (String, Maybe SBS.ByteString) ())
prepareBenchmarkProgram concurrency opts program cases = do
  let futhark = compFuthark opts

  ref_res <- runExceptT $ ensureReferenceOutput concurrency (FutharkExe futhark) "c" program cases
  case ref_res of
    Left err ->
      pure $
        Left
          ( "Reference output generation for "
              <> program
              <> " failed:\n"
              <> T.unpack err,
            Nothing
          )
    Right () -> do
      (futcode, _, futerr) <-
        liftIO $
          readProcessWithExitCode
            futhark
            ( [compBackend opts, program, "-o", binaryName program, "--server"]
                <> compOptions opts
            )
            ""

      case futcode of
        ExitSuccess -> pure $ Right ()
        ExitFailure 127 -> pure $ Left (progNotFound futhark, Nothing)
        ExitFailure _ -> pure $ Left ("Compilation of " ++ program ++ " failed:\n", Just futerr)