Starlark: Benchmarks --iterations <count> command line flag

Option to run a fixed number of iterations to use benchmarks with
an external comparison tool.

This is a follow-up to
[this comment](#12498 (comment)).

`Benchmarks --iterations 10000` will run 10000 iterations of a given
benchmark (or of the benchmarks if filter is not specified).

For example:

```
$ absh \
    -a "java -jar $HOME/Benchmarks_deploy-a.jar --filter bench_sort_small --iterations 10000000" \
    -b "java -jar $HOME/Benchmarks_deploy-b.jar --filter bench_sort_small --iterations 10000000"
...
A: N=53 r=3.992+-0.425 se=0.058
B: N=53 r=3.919+-0.384 se=0.053
B/A: 0.982
```

Note that currently `Benchmarks` parses all benchmark files (creating
unnecessary overhead), this will be fixed when new
`--filter <file>:<benchmark>` syntax is introduced.

Closes #12511.

PiperOrigin-RevId: 346108443

src/test/java/net/starlark/java/eval/Benchmarks.java

@@ -14,6 +14,7 @@
 
 package net.starlark.java.eval;
 
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.sun.management.ThreadMXBean;
 import java.io.File;
@@ -61,9 +62,9 @@
 public final class Benchmarks {
 
   private static final String HELP =
-      "Usage: Benchmarks [--help] [--filter regex] [--seconds float]\n"
-          + "Runs Starlark benchmarks matching the filter for the specified (approximate) time,\n"
-          + "and reports various performance measures.\n"
+      "Usage: Benchmarks [--help] [--filter regex] [--seconds float] [--iterations count]\n"
+          + "Runs Starlark benchmarks matching the filter for the specified approximate time or\n"
+          + "specified number of iterations, and reports various performance measures.\n"
           + "The optional filter is a regular expression applied to the string FILE:FUNC,\n"
           + "where FILE is the base name of the file and FUNC is the name of the function,\n"
           + "for example 'bench_int.star:bench_add32'.\n";
@@ -72,7 +73,8 @@ public final class Benchmarks {
 
   public static void main(String[] args) throws Exception {
     Pattern filter = null; // default: all
-    long budgetNanos = 1_000_000_000;
+    long budgetNanos = -1;
+    int iterations = -1;
 
     // parse flags
     int i;
@@ -108,6 +110,19 @@ public static void main(String[] args) throws Exception {
           fail("--seconds out of range");
         }
 
+      } else if (args[i].equals("--iterations")) {
+        if (++i == args.length) {
+          fail("--iterations needs an integer argument");
+        }
+        try {
+          iterations = Integer.parseInt(args[i]);
+        } catch (NumberFormatException e) {
+          fail("for --iterations, got '%s', want an integer number of iterations", args[i]);
+        }
+        if (iterations < 0) {
+          fail("--iterations out of range");
+        }
+
       } else {
         fail("unknown flag: %s", args[i]);
       }
@@ -116,6 +131,13 @@ public static void main(String[] args) throws Exception {
       fail("unexpected arguments");
     }
 
+    if (iterations >= 0 && budgetNanos >= 0) {
+      fail("cannot specify both --seconds and --iterations");
+    }
+    if (iterations < 0 && budgetNanos < 0) {
+      budgetNanos = 1_000_000_000;
+    }
+
     // Read testdata/bench_* files.
     File src = new File("third_party/bazel/src"); // blaze
     if (!src.exists()) {
@@ -185,8 +207,8 @@ public static void main(String[] args) throws Exception {
       for (Map.Entry<String, StarlarkFunction> e : benchmarks.entrySet()) {
         String name = e.getKey();
         System.out.flush(); // help user identify a slow benchmark
-        Benchmark b = run(name, e.getValue(), budgetNanos);
-        if (b == null) {
+        Benchmark b = new Benchmark(name, e.getValue());
+        if (!run(b, budgetNanos, iterations)) {
           ok = false;
           continue;
         }
@@ -213,50 +235,48 @@ private static void fail(String format, Object... args) {
   }
 
   // Runs benchmark function f for the specified time budget
-  // (which we may exceed by a factor of two).
-  private static Benchmark run(String name, StarlarkFunction f, long budgetNanos) {
+  // (which we may exceed by a factor of two) or number of iterations,
+  // exactly one of which must be nonnegative. Reports success.
+  private static boolean run(Benchmark b, long budgetNanos, int iterations) {
+    // Exactly one of the parameters must be specified.
+    Preconditions.checkState((budgetNanos >= 0) != (iterations >= 0));
+
     Mutability mu = Mutability.create("test");
     StarlarkThread thread = new StarlarkThread(mu, semantics);
 
-    Benchmark b = new Benchmark();
+    // Run for a fixed number of iterations?
+    if (iterations >= 0) {
+      return b.runIterations(thread, iterations);
+    }
 
-    // Keep doubling the number of iterations until we exceed the deadline.
-    // TODO(adonovan): opt: extrapolate and predict the number of iterations
-    // in the remaining time budget, being wary of extrapolation error.
-    for (b.n = 1; b.time < budgetNanos; b.n <<= 1) {
-      if (b.n <= 0) {
-        System.err.printf(
-            "In %s: function is too fast; likely a loop over `range(b.n)` is missing\n", name);
-        return null;
+    // Run for a fixed amount of time (default behavior).
+    iterations = 1;
+    while (b.time < budgetNanos) {
+      if (!b.runIterations(thread, iterations)) {
+        return false;
       }
 
-      try {
-        b.start(thread);
-        Starlark.fastcall(thread, f, new Object[] {b}, new Object[0]);
-        b.stop(thread);
-
-      } catch (EvalException ex) {
-        System.err.println(ex.getMessageWithStack());
-        return null;
-
-      } catch (
-          @SuppressWarnings("InterruptedExceptionSwallowed")
-          Throwable ex) {
-        // unhandled exception (incl. InterruptedException)
-        System.err.printf("In %s: %s\n", name, ex.getMessage());
-        ex.printStackTrace();
-        return null;
+      // Keep doubling the number of iterations until we exceed the deadline.
+      // TODO(adonovan): opt: extrapolate and predict the number of iterations
+      // in the remaining time budget, being wary of extrapolation error.
+      iterations <<= 1;
+      if (iterations <= 0) { // overflow
+        System.err.printf(
+            "In %s: function is too fast; likely a loop over `range(b.n)` is missing\n", b.name);
+        return false;
       }
     }
-
-    return b;
+    return true;
   }
 
   // The type of the parameter to each bench(b) function.
   // Provides n, the number of iterations.
   @StarlarkBuiltin(name = "Benchmark")
   private static class Benchmark implements StarlarkValue {
 
+    private final String name;
+    private final StarlarkFunction f;
+
     // The cast assumes we use the "Sun" JVM, which measures per-thread allocation and CPU.
     private final ThreadMXBean threadMX = (ThreadMXBean) ManagementFactory.getThreadMXBean();
 
@@ -276,6 +296,34 @@ private static class Benchmark implements StarlarkValue {
     private long time; // wall time (ns)
     private long steps; // Starlark computation steps
 
+    private Benchmark(String name, StarlarkFunction f) {
+      this.name = name;
+      this.f = f;
+    }
+
+    // Runs n iterations of this benchmark and reports success.
+    private boolean runIterations(StarlarkThread thread, int n) {
+      this.n = n;
+      try {
+        start(thread);
+        Starlark.fastcall(thread, f, new Object[] {this}, new Object[0]);
+        stop(thread);
+
+      } catch (EvalException ex) {
+        System.err.println(ex.getMessageWithStack());
+        return false;
+
+      } catch (
+          @SuppressWarnings("InterruptedExceptionSwallowed")
+          Throwable ex) {
+        // unhandled exception (incl. InterruptedException)
+        System.err.printf("In %s: %s\n", name, ex.getMessage());
+        ex.printStackTrace();
+        return false;
+      }
+      return true;
+    }
+
     @StarlarkMethod(name = "n", doc = "Requested number of iterations.", structField = true)
     public int n() {
       return n;