cmd/bench: add -pgo flag

This change adds a flag to cmd/bench to do a few runs to collect
profiles and feed that back into the build before doing measured
benchmark runs.

This mode is only supported in Sweet to begin with.

Change-Id: I483ca5d5385b341dee4038641e16e02f2ce3493d
Reviewed-on: https://go-review.googlesource.com/c/benchmarks/+/594776
Auto-Submit: Michael Knyszek <[email protected]>
Reviewed-by: Michael Pratt <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: mknyszek

cmd/bench/bent.go

@@ -63,7 +63,11 @@ func removeAllIncludingReadonly(dir string) error {
 	return os.RemoveAll(dir)
 }
 
-func bent(tcs []*toolchain) (err error) {
+func bent(tcs []*toolchain, pgo bool) (err error) {
+	if pgo {
+		log.Printf("Skipping bent benchmarks (PGO not supported)")
+		return nil
+	}
 	dir, err := os.MkdirTemp("", "bent")
 	if err != nil {
 		return fmt.Errorf("error creating temporary directory: %w", err)

cmd/bench/gotest.go

@@ -10,7 +10,11 @@ import (
 	"path/filepath"
 )
 
-func goTest(tcs []*toolchain) error {
+func goTest(tcs []*toolchain, pgo bool) error {
+	if pgo {
+		log.Printf("Skipping Go test benchmarks (PGO not supported)")
+		return nil
+	}
 	for _, tc := range tcs {
 		log.Printf("Running Go test benchmarks for %s", tc.Name)
 		fmt.Printf("toolchain: %s\n", tc.Name)

cmd/bench/main.go

@@ -24,6 +24,7 @@ import (
 
 var (
 	wait              = flag.Bool("wait", true, "wait for system idle before starting benchmarking")
+	pgo               = flag.Bool("pgo", false, "run the benchmarks to collect profiles and rebuild them with PGO enabled before measuring")
 	gorootExperiment  = flag.String("goroot", "", "GOROOT to test (default $GOROOT or 'go env GOROOT')")
 	gorootBaseline    = flag.String("goroot-baseline", "", "baseline GOROOT to test against (optional) (default $BENCH_BASELINE_GOROOT)")
 	branch            = flag.String("branch", "", "branch of the commits we're testing against (default $BENCH_BRANCH or unknown)")
@@ -64,16 +65,17 @@ func toolchainFromGOROOT(name, goroot string) *toolchain {
 	}
 }
 
-func run(tcs []*toolchain) error {
+func run(tcs []*toolchain, pgo bool) error {
 	// Because each of the functions below is responsible for running
 	// benchmarks under each toolchain itself, it is also responsible
 	// for ensuring that the benchmark tag "toolchain" is printed.
 	pass := true
-	if err := goTest(tcs); err != nil {
+
+	if err := goTest(tcs, pgo); err != nil {
 		pass = false
 		log.Printf("Error running Go tests: %v", err)
 	}
-	if err := bent(tcs); err != nil {
+	if err := bent(tcs, pgo); err != nil {
 		pass = false
 		log.Printf("Error running bent: %v", err)
 	}
@@ -93,7 +95,7 @@ func run(tcs []*toolchain) error {
 			return fmt.Errorf("failed to clean Go cache: %w", err)
 		}
 	}
-	if err := sweet(tcs); err != nil {
+	if err := sweet(tcs, pgo); err != nil {
 		pass = false
 		log.Printf("Error running sweet: %v", err)
 	}
@@ -189,7 +191,7 @@ func main() {
 		return
 	}
 	// Run benchmarks against the toolchains.
-	if err := run(toolchains); err != nil {
+	if err := run(toolchains, *pgo); err != nil {
 		log.Print("FAIL")
 		os.Exit(1)
 	}

cmd/bench/sweet.go

@@ -39,7 +39,7 @@ func writeSweetConfiguration(filename string, tcs []*toolchain) error {
 	return nil
 }
 
-func sweet(tcs []*toolchain) (err error) {
+func sweet(tcs []*toolchain, pgo bool) (err error) {
 	tmpDir, err := os.MkdirTemp("", "go-sweet")
 	if err != nil {
 		return fmt.Errorf("error creating temporary directory: %w", err)
@@ -103,16 +103,20 @@ func sweet(tcs []*toolchain) (err error) {
 	// Finally we can actually run the benchmarks.
 	resultsDir := filepath.Join(tmpDir, "results")
 	workDir := filepath.Join(tmpDir, "work")
-	cmd = exec.Command(
-		sweetBin, "run",
+	sweetRunArgs := []string{
+		"run",
 		"-run", "default",
 		"-count", "10",
 		"-bench-dir", filepath.Join(sweetRoot, "benchmarks"),
 		"-cache", assetsCacheDir,
 		"-work-dir", workDir,
 		"-results", resultsDir,
-		confFile,
-	)
+	}
+	if pgo {
+		sweetRunArgs = append(sweetRunArgs, "-pgo")
+	}
+	sweetRunArgs = append(sweetRunArgs, confFile)
+	cmd = exec.Command(sweetBin, sweetRunArgs...)
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
 
@@ -129,29 +133,53 @@ func sweet(tcs []*toolchain) (err error) {
 		}
 	}()
 
-	// Dump results to stdout.
+	dumpResults := func(filename string) error {
+		// Print pkg and shortname tags because Sweet won't do it.
+		benchName := filepath.Base(filepath.Dir(filename))
+		fmt.Printf("pkg: golang.org/x/benchmarks/sweet/benchmarks/%s\n", benchName)
+		fmt.Printf("shortname: sweet_%s\n", strings.ReplaceAll(benchName, "-", "_"))
+
+		// Dump results file.
+		f, err := os.Open(filename)
+		if err != nil {
+			return fmt.Errorf("opening result %s: %v", filename, err)
+		}
+		defer f.Close()
+		if _, err := io.Copy(os.Stdout, f); err != nil {
+			return fmt.Errorf("reading result %s: %v", filename, err)
+		}
+		return nil
+	}
+
+	// Dump non-PGO results to stdout.
+	fmt.Printf("pgo: off\n")
 	for _, tc := range tcs {
 		matches, err := filepath.Glob(filepath.Join(resultsDir, "*", fmt.Sprintf("%s.results", tc.Name)))
 		if err != nil {
 			return fmt.Errorf("searching for results for %s in %s: %v", tc.Name, resultsDir, err)
 		}
 		fmt.Printf("toolchain: %s\n", tc.Name)
 		for _, match := range matches {
-			// Print pkg and shortname tags because Sweet won't do it.
-			benchName := filepath.Base(filepath.Dir(match))
-			fmt.Printf("pkg: golang.org/x/benchmarks/sweet/benchmarks/%s\n", benchName)
-			fmt.Printf("shortname: sweet_%s\n", strings.ReplaceAll(benchName, "-", "_"))
+			if err := dumpResults(match); err != nil {
+				return err
+			}
+		}
+	}
 
-			// Dump results file.
-			f, err := os.Open(match)
+	// Dump PGO results to stdout, if we expect them to exist.
+	if pgo {
+		fmt.Printf("pgo: on\n")
+		for _, tc := range tcs {
+			matches, err := filepath.Glob(filepath.Join(resultsDir, "*", fmt.Sprintf("%s.pgo.results", tc.Name)))
 			if err != nil {
-				return fmt.Errorf("opening result %s: %v", match, err)
+				return fmt.Errorf("searching for PGO results for %s in %s: %v", tc.Name, resultsDir, err)
 			}
-			if _, err := io.Copy(os.Stdout, f); err != nil {
-				f.Close()
-				return fmt.Errorf("reading result %s: %v", match, err)
+			fmt.Printf("toolchain: %s\n", tc.Name)
+			for _, match := range matches {
+				if err := dumpResults(match); err != nil {
+					return err
+				}
 			}
-			f.Close()
 		}
 	}
 	return nil