From c777f305c7b8e2550c8d34c0a918aa74f9ec42bf Mon Sep 17 00:00:00 2001
From: WilliamDue <79925124+WilliamDue@users.noreply.github.com>
Date: Wed, 1 Feb 2023 17:08:32 +0100
Subject: [PATCH] `cmp-bench-json.py` rewritten in Haskell (Issue #748) (#1860)

---
 docs/man/futhark.rst        |   7 +
 futhark.cabal               |   1 +
 src/Futhark/CLI/Benchcmp.hs | 327 ++++++++++++++++++++++++++++++++++++
 src/Futhark/CLI/Main.hs     |   4 +-
 4 files changed, 338 insertions(+), 1 deletion(-)
 create mode 100644 src/Futhark/CLI/Benchcmp.hs

diff --git a/docs/man/futhark.rst b/docs/man/futhark.rst
index fceed015a5..b016343e17 100644
--- a/docs/man/futhark.rst
+++ b/docs/man/futhark.rst
@@ -24,6 +24,13 @@ documented below.
 COMMANDS
 ========
 
+futhark benchcmp FILE_A FILE_B
+------------------------------
+
+Compares two Futhark benchmarks and reports changes in performance.
+The files must be formatted in the same manner as a JSON file returned
+from :ref:`futhark-bench(1)`.
+
 futhark check [-w] PROGRAM
 --------------------------
 
diff --git a/futhark.cabal b/futhark.cabal
index 0ec50f953c..6e548408f0 100644
--- a/futhark.cabal
+++ b/futhark.cabal
@@ -184,6 +184,7 @@ library
       Futhark.CLI.C
       Futhark.CLI.CUDA
       Futhark.CLI.Check
+      Futhark.CLI.Benchcmp
       Futhark.CLI.Datacmp
       Futhark.CLI.Dataset
       Futhark.CLI.Defs
diff --git a/src/Futhark/CLI/Benchcmp.hs b/src/Futhark/CLI/Benchcmp.hs
new file mode 100644
index 0000000000..9faf4d0682
--- /dev/null
+++ b/src/Futhark/CLI/Benchcmp.hs
@@ -0,0 +1,327 @@
+-- | @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.Text qualified as T
+import Data.Vector qualified as V
+import Futhark.Bench
+import Futhark.Util (showText)
+import Futhark.Util.Options (mainWithOptions)
+import Statistics.Sample qualified as S
+import System.Console.ANSI (hSupportsANSI)
+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
+
+-- | 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 ->
+  M.Map T.Text (M.Map T.Text SpeedUp) ->
+  ([T.Text], [T.Text]) ->
+  IO ()
+printComparisons colors speedups (errors, missing) = do
+  mapM_ (putStrLn . T.unpack) $ L.sort missing
+  mapM_ (putStrLn . T.unpack) $ L.sort errors
+  mapM_ (uncurry (printProgSpeedUps colors)) $ M.toList speedups
+
+-- | Run @futhark benchcmp@
+main :: String -> [String] -> IO ()
+main = mainWithOptions () [] "<file> <file>" f
+  where
+    f [a_path', b_path'] () = 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) .)
+              . 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
diff --git a/src/Futhark/CLI/Main.hs b/src/Futhark/CLI/Main.hs
index ee29016642..b7e6f4c5e2 100644
--- a/src/Futhark/CLI/Main.hs
+++ b/src/Futhark/CLI/Main.hs
@@ -7,6 +7,7 @@ import Data.Maybe
 import Data.Text.IO qualified as T
 import Futhark.CLI.Autotune qualified as Autotune
 import Futhark.CLI.Bench qualified as Bench
+import Futhark.CLI.Benchcmp qualified as Benchcmp
 import Futhark.CLI.C qualified as C
 import Futhark.CLI.CUDA qualified as CCUDA
 import Futhark.CLI.Check qualified as Check
@@ -77,7 +78,8 @@ commands =
       ("literate", (Literate.main, "Process a literate Futhark program.")),
       ("lsp", (LSP.main, "Run LSP server.")),
       ("thanks", (Misc.mainThanks, "Express gratitude.")),
-      ("tokens", (Misc.mainTokens, "Print tokens from Futhark file."))
+      ("tokens", (Misc.mainTokens, "Print tokens from Futhark file.")),
+      ("benchcmp", (Benchcmp.main, "Compare two Futhark benchmarks."))
     ]
 
 msg :: String