Benchcmp.hs

-- | @futhark benchcmp@
module Futhark.CLI.Benchcmp (main) where

import Control.Exception (catch)
import Data.Bifunctor (Bifunctor (bimap, first, second))
import Data.ByteString.Lazy.Char8 qualified as LBS
import Data.Either qualified as E
import Data.List qualified as L
import Data.Map qualified as M
import Data.Ord (comparing)
import Data.Text qualified as T
import Data.Vector qualified as V
import Futhark.Bench
import Futhark.Util (showText)
import Futhark.Util.Options (mainWithOptions, optionsError)
import Statistics.Sample qualified as S
import System.Console.ANSI (hSupportsANSI)
import System.Console.GetOpt (ArgDescr (ReqArg), OptDescr (Option))
import System.IO (stdout)
import Text.Printf (printf)

-- | Record that summerizes a comparison between two benchmarks.
data SpeedUp = SpeedUp
  { -- | What factor the benchmark is improved by.
    speedup :: Double,
    -- | Memory usage.
    memoryUsage :: M.Map T.Text Double,
    -- | If the speedup was significant.
    significant :: Bool
  }
  deriving (Show)

-- | Terminal colors used when printing the comparisons. Some of these are not
-- colors ways of emphasising text.
data Colors = Colors
  { -- | The header color.
    header :: T.Text,
    -- | Okay color
    okblue :: T.Text,
    -- | A second okay color
    okgreen :: T.Text,
    -- | Warning color.
    warning :: T.Text,
    -- | When something fails.
    failing :: T.Text,
    -- | Default color.
    endc :: T.Text,
    -- | Bold text.
    bold :: T.Text,
    -- | Underline text.
    underline :: T.Text
  }

-- | Colors to use for a terminal device.
ttyColors :: Colors
ttyColors =
  Colors
    { header = "\ESC[95m",
      okblue = "\ESC[94m",
      okgreen = "\ESC[92m",
      warning = "\ESC[93m",
      failing = "\ESC[91m",
      endc = "\ESC[0m",
      bold = "\ESC[1m",
      underline = "\ESC[4m"
    }

-- | Colors to use for a non-terminal device.
nonTtyColors :: Colors
nonTtyColors =
  Colors
    { header = "",
      okblue = "",
      okgreen = "",
      warning = "",
      failing = "",
      endc = "",
      bold = "",
      underline = ""
    }

-- | Reads a file without throwing an error.
readFileSafely :: T.Text -> IO (Either T.Text LBS.ByteString)
readFileSafely filepath =
  (Right <$> LBS.readFile (T.unpack filepath)) `catch` couldNotRead
  where
    couldNotRead e = pure $ Left $ showText (e :: IOError)

-- | Converts DataResults to a Map with the text as a key.
toDataResultsMap :: [DataResult] -> M.Map T.Text (Either T.Text Result)
toDataResultsMap = M.fromList . fmap toTuple
  where
    toTuple (DataResult dataset dataResults) = (dataset, dataResults)

-- | Converts BenchResults to a Map with the file path as a key.
toBenchResultsMap ::
  [BenchResult] ->
  M.Map T.Text (M.Map T.Text (Either T.Text Result))
toBenchResultsMap = M.fromList . fmap toTuple
  where
    toTuple (BenchResult path dataResults) =
      (T.pack path, toDataResultsMap dataResults)

-- | Given a file path to a json file which has the form of a futhark benchmark
-- result, it will try to parse the file to a Map of Maps. The key
-- in the outer most dictionary is a file path the inner key is the dataset.
decodeFileBenchResultsMap ::
  T.Text ->
  IO (Either T.Text (M.Map T.Text (M.Map T.Text (Either T.Text Result))))
decodeFileBenchResultsMap path = do
  file <- readFileSafely path
  pure $ toBenchResultsMap <$> (file >>= (first T.pack . decodeBenchResults))

-- | Will return a text with an error saying there is a missing program in a
-- given result.
formatMissingProg :: T.Text -> T.Text -> T.Text -> T.Text
formatMissingProg = ((T.pack .) .) . printf "In %s but not %s: program %s"

-- | Will return a text with an error saying there is a missing dataset in a
-- given result.
formatMissingData :: T.Text -> T.Text -> T.Text -> T.Text -> T.Text
formatMissingData =
  (((T.pack .) .) .) . printf "In %s but not %s: program %s dataset %s"

-- | Will return texts that say there are a missing program.
formatManyMissingProg :: T.Text -> T.Text -> [T.Text] -> [T.Text]
formatManyMissingProg a_path b_path =
  zipWith3 formatMissingProg a_paths b_paths
  where
    a_paths = repeat a_path
    b_paths = repeat b_path

-- | Will return texts that say there are missing datasets for a program.
formatManyMissingData :: T.Text -> T.Text -> T.Text -> [T.Text] -> [T.Text]
formatManyMissingData prog a_path b_path =
  L.zipWith4 formatMissingData a_paths b_paths progs
  where
    a_paths = repeat a_path
    b_paths = repeat b_path
    progs = repeat prog

-- | Finds the keys two Maps does not have in common and returns a appropiate
-- error based on the functioned passed.
missingResults ::
  (T.Text -> T.Text -> [T.Text] -> [T.Text]) ->
  T.Text ->
  T.Text ->
  M.Map T.Text a ->
  M.Map T.Text b ->
  [T.Text]
missingResults toMissingMap a_path b_path a_results b_results = missing
  where
    a_keys = M.keys a_results
    b_keys = M.keys b_results
    a_missing = toMissingMap a_path b_path $ a_keys L.\\ b_keys
    b_missing = toMissingMap b_path a_path $ b_keys L.\\ a_keys
    missing = a_missing `L.union` b_missing

-- | Compares the memory usage of two results.
computeMemoryUsage ::
  M.Map T.Text Int ->
  M.Map T.Text Int ->
  M.Map T.Text Double
computeMemoryUsage a b = M.intersectionWith divide b $ M.filter (/= 0) a
  where
    divide x y = fromIntegral x / fromIntegral y

-- | Compares two results and thereby computes the Speed Up records.
compareResult :: Result -> Result -> SpeedUp
compareResult a b =
  SpeedUp
    { speedup = speedup',
      significant = significant',
      memoryUsage = memory_usage
    }
  where
    runResultToDouble :: RunResult -> Double
    runResultToDouble = fromIntegral . runMicroseconds
    toVector = V.fromList . (runResultToDouble <$>) . runResults
    a_memory_usage = memoryMap a
    b_memory_usage = memoryMap b
    a_run_results = toVector a
    b_run_results = toVector b
    a_std = S.stdDev a_run_results
    b_std = S.stdDev b_run_results
    a_mean = S.mean a_run_results
    b_mean = S.mean b_run_results
    diff = abs $ a_mean - b_mean
    speedup' = a_mean / b_mean
    significant' = diff > a_std / 2 + b_std / 2
    memory_usage = computeMemoryUsage a_memory_usage b_memory_usage

-- | Given two Maps containing datasets as keys and results as values, compare
-- the results and return the errors in a tuple.
compareDataResults ::
  T.Text ->
  T.Text ->
  T.Text ->
  M.Map T.Text (Either T.Text Result) ->
  M.Map T.Text (Either T.Text Result) ->
  (M.Map T.Text SpeedUp, ([T.Text], [T.Text]))
compareDataResults prog a_path b_path a_data b_data = result
  where
    formatMissing = formatManyMissingData prog
    partition = E.partitionEithers . fmap sequence . M.toList
    (a_errors, a_data') = second M.fromList $ partition a_data
    (b_errors, b_data') = second M.fromList $ partition b_data
    missing = missingResults formatMissing a_path b_path a_data' b_data'
    exists = M.intersectionWith compareResult a_data' b_data'
    errors = a_errors ++ b_errors
    result = (exists, (errors, missing))

-- | Given two Maps containing program file paths as keys and values as datasets
-- with results. Compare the results for each dataset in each program and
-- return the errors in a tuple.
compareBenchResults ::
  T.Text ->
  T.Text ->
  M.Map T.Text (M.Map T.Text (Either T.Text Result)) ->
  M.Map T.Text (M.Map T.Text (Either T.Text Result)) ->
  (M.Map T.Text (M.Map T.Text SpeedUp), ([T.Text], [T.Text]))
compareBenchResults a_path b_path a_bench b_bench = (exists, errors_missing)
  where
    missing = missingResults formatManyMissingProg a_path b_path a_bench b_bench
    result = M.intersectionWithKey auxiliary a_bench b_bench
    auxiliary prog = compareDataResults prog a_path b_path
    exists = M.filter (not . null) $ fst <$> result
    errors_missing' = bimap concat concat . unzip . M.elems $ snd <$> result
    errors_missing = second (missing ++) errors_missing'

-- | Formats memory usage such that it is human readable. If the memory usage
-- is not significant an empty text is returned.
memoryFormatter :: Colors -> T.Text -> Double -> T.Text
memoryFormatter colors key value
  | value < 0.99 = memoryFormat $ okgreen colors
  | value > 1.01 = memoryFormat $ failing colors
  | otherwise = ""
  where
    memoryFormat c = T.pack $ printf "%s%4.2fx@%s%s" c value key endc'
    endc' = endc colors

-- | Given a SpeedUp record the memory usage will be formatted to a colored
-- human readable text.
toMemoryText :: Colors -> SpeedUp -> T.Text
toMemoryText colors data_result
  | T.null memory_text = ""
  | otherwise = " (mem: " <> memory_text <> ")"
  where
    memory_text = M.foldrWithKey formatFolder "" memory
    memory = memoryUsage data_result
    formatFolder key value lst = lst <> memoryFormatter colors key value

-- | Given a text shorten it to a given length and add a suffix as the last
-- word.
shorten :: Int -> T.Text -> T.Text -> T.Text
shorten c end string
  | T.length string > c = (T.unwords . init $ T.words shortened) <> " " <> end
  | otherwise = string
  where
    end_len = T.length end
    (shortened, _) = T.splitAt (c - end_len) string

-- | Given a text add padding to the right of the text in form of spaces.
rightPadding :: Int -> T.Text -> T.Text
rightPadding c = T.pack . printf s
  where
    s = "%-" <> show c <> "s"

-- | Given a SpeedUp record print the SpeedUp in a human readable manner.
printSpeedUp :: Colors -> T.Text -> SpeedUp -> IO ()
printSpeedUp colors dataset data_result = do
  let color
        | significant data_result && speedup data_result > 1.01 = okgreen colors
        | significant data_result && speedup data_result < 0.99 = failing colors
        | otherwise = ""
  let short_dataset = rightPadding 64 . (<> ":") $ shorten 63 "[...]" dataset
  let memoryText = toMemoryText colors data_result
  let speedup' = speedup data_result
  let endc' = endc colors
  let format = "  %s%s%10.2fx%s%s"
  putStrLn $ printf format short_dataset color speedup' endc' memoryText

-- | Given a Map of SpeedUp records where the key is the program, print the
-- SpeedUp in a human readable manner.
printProgSpeedUps :: Colors -> T.Text -> M.Map T.Text SpeedUp -> IO ()
printProgSpeedUps colors prog bench_result = do
  putStrLn ""
  putStrLn $ printf "%s%s%s%s" (header colors) (bold colors) prog (endc colors)
  mapM_ (uncurry (printSpeedUp colors)) $ M.toList bench_result

-- | Geometric mean of a list of positive doubles; returns 1.0 for an empty list.
geoMean :: [Double] -> Double
geoMean [] = 1.0
geoMean xs = exp (sum (log <$> xs) / fromIntegral (length xs))

-- | Which metric to use when sorting program groups.
data SortMetric = BySignificantCount | ByGeoMeanSignificant | ByGeoMeanAll

-- | Precomputed per-group sort scores.
data GroupScore = GroupScore
  { gsSignificantCount :: Int,
    gsGeoMeanSignificant :: Double,
    gsGeoMeanAll :: Double
  }

computeGroupScore :: M.Map T.Text SpeedUp -> GroupScore
computeGroupScore datasets =
  GroupScore
    { gsSignificantCount = length sigReg,
      gsGeoMeanSignificant = geoMean (speedup <$> sigAll),
      gsGeoMeanAll = geoMean (speedup <$> all')
    }
  where
    all' = M.elems datasets
    sigAll = filter significant all'
    sigReg = filter (\s -> significant s && speedup s < 0.99) all'

-- | Whether to surface the worst or best programs first.
-- "Worst" means the most regressions / lowest speedup ratio.
data SortOrder = WorstFirst | BestFirst

scoreFor :: SortMetric -> GroupScore -> Double
scoreFor BySignificantCount gs = fromIntegral (gsSignificantCount gs)
scoreFor ByGeoMeanSignificant gs = gsGeoMeanSignificant gs
scoreFor ByGeoMeanAll gs = gsGeoMeanAll gs

-- | Sort a list of (program, datasets) pairs by the chosen metric and order.
-- Ties on the primary key fall back to geoMeanAll ascending then name alphabetically.
sortGroups ::
  SortMetric ->
  SortOrder ->
  [(T.Text, M.Map T.Text SpeedUp)] ->
  [(T.Text, M.Map T.Text SpeedUp)]
sortGroups metric order = L.sortBy cmpFull
  where
    score g = scoreFor metric (computeGroupScore (snd g))
    geoAll g = gsGeoMeanAll (computeGroupScore (snd g))
    -- For count metrics "worst first" = descending; for ratio metrics
    -- "worst first" = ascending (lower ratio = bigger regression).
    primaryCmp = case (metric, order) of
      (BySignificantCount, WorstFirst) -> flip compare `on` score
      (BySignificantCount, BestFirst) -> compare `on` score
      (_, WorstFirst) -> compare `on` score
      (_, BestFirst) -> flip compare `on` score
    cmpFull a b =
      primaryCmp a b
        <> comparing geoAll a b
        <> comparing fst a b
    on f g x y = f (g x) (g y)

-- | Config for the benchcmp tool.
data BenchcmpConfig = BenchcmpConfig
  { cfgSortMetric :: Maybe SortMetric,
    cfgSortOrder :: SortOrder
  }

defaultConfig :: BenchcmpConfig
defaultConfig = BenchcmpConfig Nothing WorstFirst

-- | Given a Map of programs with dataset speedups and relevant errors, print
-- the errors and print the speedups in a human readable manner.
printComparisons ::
  Colors ->
  BenchcmpConfig ->
  M.Map T.Text (M.Map T.Text SpeedUp) ->
  ([T.Text], [T.Text]) ->
  IO ()
printComparisons colors cfg speedups (errors, missing) = do
  mapM_ (putStrLn . T.unpack) $ L.sort missing
  mapM_ (putStrLn . T.unpack) $ L.sort errors
  let groups = case cfgSortMetric cfg of
        Nothing -> M.toList speedups
        Just metric -> sortGroups metric (cfgSortOrder cfg) (M.toList speedups)
  mapM_ (uncurry (printProgSpeedUps colors)) groups

-- | Run @futhark benchcmp@
main :: String -> [String] -> IO ()
main = mainWithOptions defaultConfig options "<file> <file>" f
  where
    options =
      [ Option
          []
          ["sort-by"]
          ( ReqArg
              ( \arg -> case arg of
                  "significant" -> Right $ \cfg -> cfg {cfgSortMetric = Just BySignificantCount}
                  "geomean-significant" -> Right $ \cfg -> cfg {cfgSortMetric = Just ByGeoMeanSignificant}
                  "geomean-all" -> Right $ \cfg -> cfg {cfgSortMetric = Just ByGeoMeanAll}
                  _ -> Left . optionsError $ "Unknown --sort-by value: " <> arg
              )
              "METRIC"
          )
          "Sort program groups by: significant, geomean-significant, geomean-all (default: unsorted).",
        Option
          []
          ["order"]
          ( ReqArg
              ( \arg -> case arg of
                  "worst-first" -> Right $ \cfg -> cfg {cfgSortOrder = WorstFirst}
                  "best-first" -> Right $ \cfg -> cfg {cfgSortOrder = BestFirst}
                  _ -> Left . optionsError $ "Unknown --order value: " <> arg
              )
              "ORDER"
          )
          "Sort order: worst-first (default) or best-first."
      ]
    f [a_path', b_path'] cfg = Just $ do
      let a_path = T.pack a_path'
      let b_path = T.pack b_path'
      a_either <- decodeFileBenchResultsMap a_path
      b_either <- decodeFileBenchResultsMap b_path

      isTty <- hSupportsANSI stdout

      let colors =
            if isTty
              then ttyColors
              else nonTtyColors

      let comparePrint =
            (uncurry (printComparisons colors cfg) .)
              . compareBenchResults a_path b_path

      case (a_either, b_either) of
        (Left a, Left b) -> putStrLn . T.unpack $ (a <> "\n" <> b)
        (Left a, _) -> putStrLn . T.unpack $ a
        (_, Left b) -> putStrLn . T.unpack $ b
        (Right a, Right b) -> comparePrint a b
    f _ _ = Nothing