Lightweight benchmarking in the uplc excutable

plutus-executables/executables/uplc/Main.hs

@@ -6,11 +6,22 @@ module Main (main) where
 import PlutusCore qualified as PLC
 import PlutusCore.Annotation (SrcSpan)
 import PlutusCore.Data (Data)
-import PlutusCore.Default (BuiltinSemanticsVariant (..))
-import PlutusCore.Evaluation.Machine.ExBudget (ExBudget (..), ExRestrictingBudget (..))
+import PlutusCore.Default
+  ( BuiltinSemanticsVariant (..)
+  )
+import PlutusCore.Evaluation.Machine.ExBudget
+  ( ExBudget (..)
+  , ExRestrictingBudget (..)
+  )
 import PlutusCore.Evaluation.Machine.ExBudgetingDefaults qualified as PLC
-import PlutusCore.Evaluation.Machine.ExMemory (ExCPU (..), ExMemory (..))
-import PlutusCore.Executable.AstIO (UplcTermNDB, toDeBruijnTermUPLC)
+import PlutusCore.Evaluation.Machine.ExMemory
+  ( ExCPU (..)
+  , ExMemory (..)
+  )
+import PlutusCore.Executable.AstIO
+  ( UplcTermNDB
+  , toDeBruijnTermUPLC
+  )
 import PlutusCore.Executable.Blueprint
 import PlutusCore.Executable.Common
 import PlutusCore.Executable.Eval
@@ -24,30 +35,57 @@ import UntypedPlutusCore.Evaluation.Machine.Cek qualified as Cek
 import UntypedPlutusCore.Evaluation.Machine.SteppableCek.DebugDriver qualified as D
 import UntypedPlutusCore.Evaluation.Machine.SteppableCek.Internal qualified as D
 
-import Codec.Serialise (DeserialiseFailure, deserialiseOrFail)
+import Codec.Serialise
+  ( DeserialiseFailure
+  , deserialiseOrFail
+  )
 import Control.DeepSeq (force)
-import Control.Monad.Except (runExcept, tryError)
+import Control.Exception (evaluate)
+import Control.Monad.Except
+  ( runExcept
+  , tryError
+  )
 import Control.Monad.Extra (whenJust)
 import Control.Monad.IO.Class (liftIO)
 import Control.Monad.ST (RealWorld)
-import Criterion (benchmarkWith, whnf)
+import Criterion
+  ( benchmarkWith
+  , whnf
+  )
 import Criterion.Main (defaultConfig)
 import Criterion.Types (Config (..))
 import Data.ByteString.Base16 qualified as Base16
 import Data.ByteString.Lazy qualified as BSL
+import Data.IORef
+  ( newIORef
+  , readIORef
+  )
 import Data.List.Split (splitOn)
 import Data.Text qualified as T
 import Data.Text.Encoding qualified as T
 import Data.Text.IO qualified as T
-import Data.Time.Clock.System (getSystemTime, systemNanoseconds)
+import Data.Time.Clock.System
+  ( getSystemTime
+  , systemNanoseconds
+  )
 import Options.Applicative
 import PlutusCore.Flat (unflat)
 import Prettyprinter ((<+>))
+import System.CPUTime (getCPUTime)
 import System.Console.Haskeline qualified as Repl
 import System.Directory.Extra (doesFileExist)
-import System.Exit (ExitCode (..), exitFailure, exitWith)
+import System.Exit
+  ( ExitCode (..)
+  , exitFailure
+  , exitWith
+  )
 import System.FilePath
-import System.IO (hPrint, stderr)
+import System.IO
+  ( hPrint
+  , stderr
+  )
+import System.Mem (performGC)
+import Text.Printf (printf)
 import Text.Read (readMaybe)
 
 import Data.Version.Extras (gitAwareVersionInfo)
@@ -82,6 +120,14 @@ data EvalOptions
       CekModel
       (BuiltinSemanticsVariant PLC.DefaultFun)
 
+data TimeEvalOptions
+  = TimeEvalOptions
+      Input
+      Format
+      (BuiltinSemanticsVariant PLC.DefaultFun)
+      Integer -- number of repetitions
+      Bool -- raw output (nanoseconds, no units)
+
 data BenchmarkOptions
   = BenchmarkOptions
       Input
@@ -108,6 +154,7 @@ data Command
   | Print PrintOptions
   | Example ExampleOptions
   | Eval EvalOptions
+  | TimeEval TimeEvalOptions
   | Dbg DbgOptions
   | DumpModel (BuiltinSemanticsVariant PLC.DefaultFun)
   | PrintBuiltinSignatures
@@ -153,6 +200,26 @@ evalOpts =
     <*> cekmodel
     <*> builtinSemanticsVariant
 
+timeOpts :: Parser TimeEvalOptions
+timeOpts =
+  TimeEvalOptions
+    <$> input
+    <*> inputformat
+    <*> builtinSemanticsVariant
+    <*> option
+      auto
+      ( short 'n'
+          <> long "repeat"
+          <> metavar "N"
+          <> value 100
+          <> showDefault
+          <> help "Number of times to evaluate the script (average time is reported for N > 1)."
+      )
+    <*> switch
+      ( long "raw"
+          <> help "Print the average time in nanoseconds with no units."
+      )
+
 dbgOpts :: Parser DbgOptions
 dbgOpts =
   DbgOptions
@@ -312,6 +379,15 @@ plutusOpts =
             (Eval <$> evalOpts)
             (progDesc "Evaluate an untyped Plutus Core program using the CEK machine.")
         )
+      <> command
+        "time"
+        ( info
+            (TimeEval <$> timeOpts)
+            ( progDesc $
+                "Time the evaluation of an untyped Plutus Core program using the CEK machine.  "
+                  ++ "For best results, bypass cabal and run the `uplc` binary directly; for example use `$(cabal list-bin uplc)`."
+            )
+        )
       <> command
         "debug"
         ( info
@@ -594,24 +670,33 @@ runBenchmark :: BenchmarkOptions -> IO ()
 runBenchmark (BenchmarkOptions inp ifmt semvar timeLim) = do
   prog <- readProgram ifmt inp
   let criterionConfig = defaultConfig {reportFile = Nothing, timeLimit = timeLim}
-      cekparams = PLC.defaultCekParametersForVariant semvar
-      -- Extract an evaluation result
-      getResult = either (error . show) (const ()) . Cek.cekResultToEither . Cek._cekReportResult
-      evaluate = getResult . Cek.runCekDeBruijn cekparams Cek.restrictingEnormous Cek.noEmitter
+      evalCtx = mkDefaultEvalCtx semvar
+      -- Evaluate the term the same way the 'time' subcommand (and production)
+      -- does, erroring on an unexpected failure.
+      cekEval = either (error . show) (const ()) . evaluateCekLikeInProd evalCtx
       -- readProgam throws away De Bruijn indices and returns an AST with Names;
       -- we have to put them back to get an AST with NamedDeBruijn names.
-      !term =
+      term =
         fromRight (error "Unexpected open term in runBenchmark.")
           . runExcept @FreeVariableError
           $ UPLC.deBruijnTerm (UPLC._progTerm prog)
       -- Big names slow things down
-      !anonTerm = UPLC.termMapNames (\(PLC.NamedDeBruijn _ i) -> PLC.NamedDeBruijn "" i) term
-      -- Big annotations slow things down
+      anonTerm = UPLC.termMapNames (\(PLC.NamedDeBruijn _ i) -> PLC.NamedDeBruijn "" i) term
+      -- Big annotations slow things down.  Forcing this to NF here drives all
+      -- the preparatory work (de Bruijn conversion, name anonymisation, 'void')
+      -- to completion before benchmarking, so it isn't measured by Criterion.
       !unitAnnTerm = force (void anonTerm)
-  benchmarkWith criterionConfig $! whnf evaluate unitAnnTerm
+  benchmarkWith criterionConfig $! whnf cekEval unitAnnTerm
 
 ---------------- Evaluation ----------------
 
+formatNs :: Integer -> String
+formatNs ns
+  | ns < 1000 = printf "%d ns" ns
+  | ns < 1000000 = printf "%.3f \xb5s" (fromIntegral ns / 1000 :: Double)
+  | ns < 1000000000 = printf "%.3f ms" (fromIntegral ns / 1000000 :: Double)
+  | otherwise = printf "%.3f s" (fromIntegral ns / 1000000000 :: Double)
+
 runEval :: EvalOptions -> IO ()
 runEval
   ( EvalOptions
@@ -645,24 +730,74 @@ runEval
           Silent -> SomeBudgetMode Cek.restrictingEnormous
           Verbose bm -> bm
     case budgetM of
-      SomeBudgetMode bm ->
-        do
-          let Cek.CekReport res budget logs = Cek.runCek cekparams bm emitM term
-          case Cek.cekResultToEither res of
-            Left err -> hPrint stderr err
-            Right v ->
-              case nameFormat of
-                IdNames -> writeToOutput outp $ prettyPrintByMode printMode v
-                DeBruijnNames -> writeToOutput outp $ prettyPrintByMode printMode $ toDeBruijnTermUPLC v
-          case budgetMode of
-            Silent -> pure ()
-            Verbose _ -> printBudgetState term cekModel budget
-          case traceMode of
-            None -> pure ()
-            _ -> writeToOutput outp (T.intercalate "\n" logs)
-          case Cek.cekResultToEither res of
-            Left _ -> exitFailure
-            Right _ -> pure ()
+      SomeBudgetMode bm -> do
+        report <- evaluate (Cek.runCek cekparams bm emitM term)
+        let Cek.CekReport res budget logs = report
+        case Cek.cekResultToEither res of
+          Left err -> hPrint stderr err
+          Right v ->
+            case nameFormat of
+              IdNames -> writeToOutput outp $ prettyPrintByMode printMode v
+              DeBruijnNames -> writeToOutput outp $ prettyPrintByMode printMode $ toDeBruijnTermUPLC v
+        case budgetMode of
+          Silent -> pure ()
+          Verbose _ -> printBudgetState term cekModel budget
+        case traceMode of
+          None -> pure ()
+          _ -> writeToOutput outp (T.intercalate "\n" logs)
+        case Cek.cekResultToEither res of
+          Left _ -> exitFailure
+          Right _ -> pure ()
+
+---------------- Timing ----------------
+
+runTimeEval :: TimeEvalOptions -> IO ()
+runTimeEval (TimeEvalOptions inp ifmt semvar n raw) = do
+  prog <- readProgram ifmt inp
+  let count = fromIntegral (max 1 n) :: Int
+      term = void $ prog ^. UPLC.progTerm
+      dbTerm =
+        fromRight (error "time: term has free variables")
+          . runExcept @FreeVariableError
+          $ UPLC.deBruijnTerm term
+  -- Fully evaluate the anonymised term to NF before timing.  This single
+  -- 'force' drives all the preparatory work (the 'void', the de Bruijn
+  -- conversion and the name anonymisation) to completion, since 'anonTerm'
+  -- depends on all of it; doing it here keeps it out of the timed loop.
+  anonTerm <-
+    evaluate . force $
+      UPLC.termMapNames (\(PLC.NamedDeBruijn _ i) -> PLC.NamedDeBruijn "" i) dbTerm
+  let !evalCtx = mkDefaultEvalCtx semvar
+  performGC
+  -- Store the term in an IORef so GHC cannot CSE/share the result of
+  -- evaluateCekLikeInProd across iterations.
+  termRef <- newIORef anonTerm
+  -- We measure CPU time (getCPUTime) rather than wall-clock time
+  -- (getSystemTime): a CPU-time clock does not advance while the thread is
+  -- descheduled, so the measurement is immune to contention from co-runners
+  -- sharing the core.  In particular, running via 'cabal run' pins the
+  -- long-lived cabal supervisor to the same core, and its periodic wake-ups
+  -- would otherwise be charged to whichever evaluation happened to be in
+  -- progress, inflating the reported time (especially for inputs that take
+  -- longer to parse and so keep the process alive longer).
+  let loop 0 lastOk !total = pure (lastOk, total)
+      loop k _ !total = do
+        term' <- readIORef termRef
+        t0 <- getCPUTime
+        r <- evaluate (evaluateCekLikeInProd evalCtx term')
+        t1 <- getCPUTime
+        let !ok = either (const False) (const True) r
+        loop (k - 1) ok (total + (t1 - t0))
+  (lastOk, totalPs) <- loop count True 0
+  -- getCPUTime returns picoseconds; convert the total to nanoseconds.
+  let avgNs = totalPs `div` (1000 * fromIntegral count)
+  putStrLn $
+    if raw
+      then show avgNs
+      else
+        (if count > 1 then "Average evaluation time (" ++ show count ++ " runs): " else "Evaluation time: ")
+          ++ formatNs avgNs
+  if lastOk then pure () else exitFailure
 
 ---------------- Debugging ----------------
 
@@ -755,6 +890,7 @@ main = do
     ApplyToCborData opts -> runApplyToCborData opts
     Benchmark opts -> runBenchmark opts
     Eval opts -> runEval opts
+    TimeEval opts -> runTimeEval opts
     Dbg opts -> runDbg opts
     Example opts -> runUplcPrintExample opts
     Optimise opts -> runOptimisations opts

plutus-ledger-api/executables/src/PlutusCore/Executable/Eval.hs

@@ -78,6 +78,22 @@ evalOptimizerTrace evalCtx trace args =
       )
         <$> allASTs trace
 
+{- TODO: This is an exact copy of some code in `PlutusBenchmark.Common`.  Check
+ if we can use this version in plutus-benchmark without affecting the
+ benchmark results (initial experiments were unclear). -}
+{-| Evaluate a term as it would be evaluated using the on-chain evaluator,
+at the most recent protocol version with restrictingEnormous budget mode
+(no budget tracking overhead). Suitable for timing. -}
+evaluateCekLikeInProd
+  :: EvaluationContext
+  -> UPLC.Term UPLC.NamedDeBruijn UPLC.DefaultUni UPLC.DefaultFun ()
+  -> Either
+       (CekEvaluationException UPLC.NamedDeBruijn UPLC.DefaultUni UPLC.DefaultFun)
+       (UPLC.Term UPLC.NamedDeBruijn UPLC.DefaultUni UPLC.DefaultFun ())
+evaluateCekLikeInProd evalCtx term =
+  cekResultToEither . _cekReportResult $
+    evaluateTerm restrictingEnormous newestPV Quiet evalCtx term
+
 {-| Evaluate a single program term applied to arguments in counting mode.
 Returns @(Maybe error, budget)@. -}
 evalCountingWithArgs