From 610fe4adaae2e773858870445dfcfc83c63c84e2 Mon Sep 17 00:00:00 2001
From: Alexey Radul <axch@google.com>
Date: Wed, 31 Aug 2022 13:38:01 -0700
Subject: [PATCH] Add a unit test mode that runs benchmarks to check that they
 do not bit-rot.

---
 benchmarks/fused_sum.dx |  5 +++++
 makefile                | 18 +++++++++++++++---
 src/lib/LLVMExec.hs     |  5 ++++-
 src/lib/PPrint.hs       |  6 ++++--
 4 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/benchmarks/fused_sum.dx b/benchmarks/fused_sum.dx
index 76de21e7e..d1048611c 100644
--- a/benchmarks/fused_sum.dx
+++ b/benchmarks/fused_sum.dx
@@ -4,3 +4,8 @@ n = if dex_test_mode() then 1000 else 1000000
 sum $ for i:(Fin n).
   x = n_to_i64 (ordinal i)
   x * x
+> 332833500
+>
+> fused-sum
+> Compile time: 22.192 ms
+> Run time:     5.186 us 	(based on 1 run)
diff --git a/makefile b/makefile
index 99454f7d7..b4ddc2ad7 100644
--- a/makefile
+++ b/makefile
@@ -216,9 +216,13 @@ test-names = uexpr-tests adt-tests type-tests eval-tests show-tests read-tests \
 
 doc-names = conditionals functions
 
-all-names = $(test-names:%=tests/%) $(example-names:%=examples/%) $(doc-names:%=doc/%)
+benchmark-names = fused_sum
 
-quine-test-targets = $(all-names:%=run-%)
+quine-test-targets = \
+  $(test-names:%=run-tests/%) \
+  $(example-names:%=run-examples/%) \
+  $(doc-names:%=run-doc/%) \
+  $(benchmark-names:%=run-bench-tests/%)
 
 update-test-targets    = $(test-names:%=update-tests/%)
 update-doc-targets     = $(doc-names:%=update-doc/%)
@@ -264,6 +268,10 @@ run-doc/%: doc/%.dx just-build
 	misc/check-quine $< $(dex) script
 run-examples/%: examples/%.dx just-build
 	misc/check-quine $< $(dex) -O script
+# This runs the benchmark in test mode, which means we're checking
+# that it's not broken, but not actually trying to measure runtimes
+run-bench-tests/%: benchmarks/%.dx just-build
+	misc/check-quine $< $(dex) -O script
 
 lower-tests: export DEX_LOWER=1
 lower-tests: export DEX_ALLOW_CONTRACTIONS=0
@@ -301,7 +309,11 @@ update-examples/%: examples/%.dx just-build
 	$(dex) script $< > $<.tmp
 	mv $<.tmp $<
 
-run-gpu-tests: export DEX_ALLOC_CONTRACTIONS=0
+update-bench-tests/%: benchmarks/%.dx just-build
+	$(dex) script $< > $<.tmp
+	mv $<.tmp $<
+
+run-gpu-tests: export DEX_ALLOW_CONTRACTIONS=0
 run-gpu-tests: tests/gpu-tests.dx just-build
 	misc/check-quine $< $(dex) --backend llvm-cuda script
 
diff --git a/src/lib/LLVMExec.hs b/src/lib/LLVMExec.hs
index 560c8fcce..18b55ea9d 100644
--- a/src/lib/LLVMExec.hs
+++ b/src/lib/LLVMExec.hs
@@ -110,8 +110,11 @@ compileAndBench shouldSyncCUDA logger objFiles ast fname args resultTypes = do
                   freeLitVals resultPtr resultTypes
             let sync = when shouldSyncCUDA $ synchronizeCUDA
             exampleDuration <- snd <$> measureSeconds (run >> sync)
+            test_mode <- (Just "t" ==) <$> E.lookupEnv "DEX_TEST_MODE"
             let timeBudget = 2 -- seconds
-            let benchRuns = (ceiling $ timeBudget / exampleDuration) :: Int
+            let benchRuns = if test_mode
+                  then 1
+                  else (ceiling $ timeBudget / exampleDuration) :: Int
             sync
             totalTime <- liftM snd $ measureSeconds $ do
               forM_ [1..benchRuns] $ const run
diff --git a/src/lib/PPrint.hs b/src/lib/PPrint.hs
index 8b5c8a23d..034ad778b 100644
--- a/src/lib/PPrint.hs
+++ b/src/lib/PPrint.hs
@@ -576,8 +576,10 @@ instance Pretty Output where
     benchName <> hardline <>
     "Compile time: " <> prettyDuration compileTime <> hardline <>
     "Run time:     " <> prettyDuration runTime <+>
-    (case stats of Just (runs, _) -> "\t" <> parens ("based on" <+> p runs <+> "runs")
-                   Nothing        -> "")
+    (case stats of
+       Just (runs, _) ->
+         "\t" <> parens ("based on" <+> p runs <+> plural "run" "runs" runs)
+       Nothing        -> "")
     where benchName = case name of "" -> ""
                                    _  -> "\n" <> p name
   pretty (PassInfo _ s) = p s