sweet: add mode for PGO benchmarking

With -pgo, for each configuration sweet automatically runs an initial
profiling configuration. The merged profiles from these runs are used as
the PGO input to a ".pgo" variant of the configuration. Comparing the
base configuration to the ".pgo" variant indicates the effect of PGO.

At a lower level, the config format adds a "pgofiles" map, which can be
used to specify PGO profiles to use for each benchmark.

Ultimately this sets GOFLAGS=-pgo=/path in BuildEnv. Some benchmarks may
not currently properly plumb this into their build (e.g., the GoBuild
benchmarks don't build the compiler at all). Existing benchmarks need to
be double-checked that they actually get PGO enabled.

For golang/go#55022.

Change-Id: I81c0cb085ef3b5196a05d2565dd0e2f83057b9fa

sweet/benchmarks/internal/cgroups/cgroups.go

@@ -110,4 +110,3 @@ func (c *Cmd) RSSFunc() func() (uint64, error) {
  return strconv.ParseUint(strings.TrimSpace(string(data)), 10, 64)
  }
 }
-

sweet/cmd/sweet/benchmark.go

@@ -202,20 +202,27 @@ func (b *benchmark) execute(cfgs []*common.Config, r *runCfg) error {
  return err
  }
 
- // Retrieve the benchmark's source.
- if err := b.harness.Get(srcDir); err != nil {
- return fmt.Errorf("retrieving source for %s: %v", b.name, err)
+ // Retrieve the benchmark's source, if needed. If execute is called
+ // multiple times, this will already be done.
+ _, err := os.Stat(srcDir)
+ if os.IsNotExist(err) {
+ if err := b.harness.Get(srcDir); err != nil {
+ return fmt.Errorf("retrieving source for %s: %v", b.name, err)
+ }
  }
 
  // Create the results directory for the benchmark.
- resultsDir := filepath.Join(r.resultsDir, b.name)
+ resultsDir := r.benchmarkResultsDir(b)
  if err := mkdirAll(resultsDir); err != nil {
  return fmt.Errorf("creating results directory for %s: %v", b.name, err)
  }
 
  // Perform a setup step for each config for the benchmark.
  setups := make([]common.RunConfig, 0, len(cfgs))
- for _, cfg := range cfgs {
+ for _, pcfg := range cfgs {
+ // Local copy for per-benchmark environment adjustments.
+ cfg := pcfg.Copy()
+
  // Create directory hierarchy for benchmarks.
  workDir := filepath.Join(topDir, cfg.Name)
  binDir := filepath.Join(workDir, "bin")
@@ -236,6 +243,16 @@ func (b *benchmark) execute(cfgs []*common.Config, r *runCfg) error {
  }
  }
 
+ // Add PGO if profile specified for this benchmark.
+ if pgo, ok := cfg.PGOFiles[b.name]; ok {
+ goflags, ok := cfg.BuildEnv.Lookup("GOFLAGS")
+ if ok {
+ goflags += " "
+ }
+ goflags += fmt.Sprintf("-pgo=%s", pgo)
+ cfg.BuildEnv.Env = cfg.BuildEnv.MustSet("GOFLAGS=" + goflags)
+ }
+
  // Build the benchmark (application and any other necessary components).
  bcfg := common.BuildConfig{
  BinDir: binDir,
@@ -264,7 +281,7 @@ func (b *benchmark) execute(cfgs []*common.Config, r *runCfg) error {
  }
  if r.cpuProfile || r.memProfile || r.perf {
  // Create a directory for any profile files to live in.
- resultsProfilesDir := filepath.Join(resultsDir, fmt.Sprintf("%s.debug", cfg.Name))
+ resultsProfilesDir := r.runProfilesDir(b, cfg)
  mkdirAll(resultsProfilesDir)
 
  // We need to pass arguments to the benchmark binary to generate

sweet/cmd/sweet/integration_test.go

@@ -20,6 +20,15 @@ import (
 )
 
 func TestSweetEndToEnd(t *testing.T) {
+ t.Run("standard", func(t *testing.T) {
+ testSweetEndToEnd(t, false)
+ })
+ t.Run("pgo", func(t *testing.T) {
+ testSweetEndToEnd(t, true)
+ })
+}
+
+func testSweetEndToEnd(t *testing.T, pgo bool) {
  if runtime.GOOS != "linux" || runtime.GOARCH != "amd64" {
  t.Skip("Sweet is currently only fully supported on linux/amd64")
  }
@@ -39,6 +48,17 @@ func TestSweetEndToEnd(t *testing.T) {
  Env: common.NewEnvFromEnviron(),
  }
 
+ if pgo {
+ cmd := exec.Command(goTool.Tool, "help", "build")
+ out, err := cmd.Output()
+ if err != nil {
+ t.Fatalf("error running go help build: %v", err)
+ }
+ if !strings.Contains(string(out), "-pgo") {
+ t.Skip("toolchain missing -pgo support")
+ }
+ }
+
  // Build sweet.
  wd, err := os.Getwd()
  if err != nil {
@@ -103,7 +123,8 @@ func TestSweetEndToEnd(t *testing.T) {
 
  var outputMu sync.Mutex
  runShard := func(shard, resultsDir, workDir string) {
- runCmd := exec.Command(sweetBin, "run",
+ args := []string{
+ "run",
  "-run", shard,
  "-shell",
  "-count", "1",
@@ -112,8 +133,12 @@ func TestSweetEndToEnd(t *testing.T) {
  "-results", resultsDir,
  "-work-dir", workDir,
  "-short",
- cfgPath,
- )
+ }
+ if pgo {
+ args = append(args, "-pgo")
+ }
+ args = append(args, cfgPath)
+ runCmd := exec.Command(sweetBin, args...)
  output, runErr := runCmd.CombinedOutput()
 
  outputMu.Lock()

sweet/cmd/sweet/run.go

@@ -12,6 +12,7 @@ import (
  "io/fs"
  "os"
  "path/filepath"
+ "regexp"
  "sort"
  "strings"
  "unicode/utf8"
@@ -21,6 +22,7 @@ import (
  "golang.org/x/benchmarks/sweet/common/log"
 
  "github.com/BurntSushi/toml"
+ "github.com/google/pprof/profile"
 )
 
 type csvFlag []string
@@ -41,6 +43,8 @@ files.`
 `
 )
 
+const pgoCountDefaultMax = 5
+
 type runCfg struct {
  count int
  resultsDir string
@@ -53,6 +57,8 @@ type runCfg struct {
  memProfile bool
  perf bool
  perfFlags string
+ pgo bool
+ pgoCount int
  short bool
 
  assetsFS fs.FS
@@ -67,6 +73,14 @@ func (r *runCfg) logCopyDirCommand(fromRelDir, toDir string) {
  }
 }
 
+func (r *runCfg) benchmarkResultsDir(b *benchmark) string {
+ return filepath.Join(r.resultsDir, b.name)
+}
+
+func (r *runCfg) runProfilesDir(b *benchmark, c *common.Config) string {
+ return filepath.Join(r.benchmarkResultsDir(b), fmt.Sprintf("%s.debug", c.Name))
+}
+
 type runCmd struct {
  runCfg
  quiet bool
@@ -135,6 +149,8 @@ func (c *runCmd) SetFlags(f *flag.FlagSet) {
  f.BoolVar(&c.runCfg.memProfile, "memprofile", false, "whether to dump a memory profile for each benchmark (ensures all executions do the same amount of work")
  f.BoolVar(&c.runCfg.perf, "perf", false, "whether to run each benchmark under Linux perf and dump the results")
  f.StringVar(&c.runCfg.perfFlags, "perf-flags", "", "the flags to pass to Linux perf if -perf is set")
+ f.BoolVar(&c.pgo, "pgo", false, "perform PGO testing; for each config, collect profiles from a baseline run which are used to feed into a generated PGO config")
+ f.IntVar(&c.runCfg.pgoCount, "pgo-count", 0, "the number of times to run profiling runs for -pgo; defaults to the value of -count if <=5, or 5 if higher")
  f.IntVar(&c.runCfg.count, "count", 25, "the number of times to run each benchmark")
 
  f.BoolVar(&c.quiet, "quiet", false, "whether to suppress activity output on stderr (no effect on -shell)")
@@ -153,6 +169,13 @@ func (c *runCmd) Run(args []string) error {
  log.SetCommandTrace(c.printCmd)
  log.SetActivityLog(!c.quiet)
 
+ if c.runCfg.pgoCount == 0 {
+ c.runCfg.pgoCount = c.runCfg.count
+ if c.runCfg.pgoCount > pgoCountDefaultMax {
+ c.runCfg.pgoCount = pgoCountDefaultMax
+ }
+ }
+
  var err error
  if c.workDir == "" {
  // Create a temporary work tree for running the benchmarks.
@@ -267,6 +290,14 @@ func (c *runCmd) Run(args []string) error {
  if config.ExecEnv.Env == nil {
  config.ExecEnv.Env = common.NewEnvFromEnviron()
  }
+ if config.PGOFiles == nil {
+ config.PGOFiles = make(map[string]string)
+ }
+ for k := range config.PGOFiles {
+ if _, ok := allBenchmarksMap[k]; !ok {
+ return fmt.Errorf("config %q in %q pgofiles references unknown benchmark %q", config.Name, configFile, k)
+ }
+ }
  configs = append(configs, config)
  }
  }
@@ -304,6 +335,14 @@ func (c *runCmd) Run(args []string) error {
  }
  }
 
+ // Collect profiles from baseline runs and create new PGO'd configs.
+ if c.pgo {
+ configs, err = c.preparePGO(configs, benchmarks)
+ if err != nil {
+ return fmt.Errorf("error preparing PGO profiles: %w", err)
+ }
+ }
+
  // Execute each benchmark for all configs.
  var errEncountered bool
  for _, b := range benchmarks {
@@ -321,6 +360,113 @@ func (c *runCmd) Run(args []string) error {
  return nil
 }
 
+func (c *runCmd) preparePGO(configs []*common.Config, benchmarks []*benchmark) ([]*common.Config, error) {
+ profileConfigs := make([]*common.Config, 0, len(configs))
+ for _, c := range configs {
+ cc := c.Copy()
+ cc.Name += ".profile"
+ profileConfigs = append(profileConfigs, cc)
+ }
+
+ profileRunCfg := c.runCfg
+ profileRunCfg.cpuProfile = true
+ profileRunCfg.count = profileRunCfg.pgoCount
+
+ log.Printf("Running profile collection runs")
+
+ // Execute benchmarks to collect profiles.
+ var errEncountered bool
+ for _, b := range benchmarks {
+ if err := b.execute(profileConfigs, &profileRunCfg); err != nil {
+ if c.stopOnError {
+ return nil, err
+ }
+ errEncountered = true
+ log.Error(err)
+ }
+ }
+ if errEncountered {
+ return nil, fmt.Errorf("failed to execute profile collection benchmarks, see log for details")
+ }
+
+ // Merge all the profiles and add new PGO configs.
+ newConfigs := configs
+ for i := range configs {
+ origConfig := configs[i]
+ profileConfig := profileConfigs[i]
+ pgoConfig := origConfig.Copy()
+ pgoConfig.Name += ".pgo"
+ pgoConfig.PGOFiles = make(map[string]string)
+
+ for _, b := range benchmarks {
+ p, err := mergeCPUProfiles(profileRunCfg.runProfilesDir(b, profileConfig))
+ if err != nil {
+ return nil, fmt.Errorf("error merging profiles for %s/%s: %w", b.name, profileConfig.Name, err)
+ }
+ pgoConfig.PGOFiles[b.name] = p
+ }
+
+ newConfigs = append(newConfigs, pgoConfig)
+ }
+
+ return newConfigs, nil
+}
+
+var cpuProfileRe = regexp.MustCompile(`^.*\.cpu[0-9]+$`)
+
+func mergeCPUProfiles(dir string) (string, error) {
+ entries, err := os.ReadDir(dir)
+ if err != nil {
+ return "", fmt.Errorf("error reading dir %q: %w", dir, err)
+ }
+
+ var profiles []*profile.Profile
+ addProfile := func(name string) error {
+ f, err := os.Open(name)
+ if err != nil {
+ return fmt.Errorf("error opening profile %q: %w", name, err)
+ }
+ defer f.Close()
+
+ p, err := profile.Parse(f)
+ if err != nil {
+ return fmt.Errorf("error parsing profile %q: %w", name, err)
+ }
+ profiles = append(profiles, p)
+ return nil
+ }
+
+ for _, e := range entries {
+ name := e.Name()
+ if cpuProfileRe.FindString(name) == "" {
+ continue
+ }
+
+ if err := addProfile(filepath.Join(dir, name)); err != nil {
+ return "", err
+ }
+ }
+
+ if len(profiles) == 0 {
+ return "", fmt.Errorf("no profiles found in %q", dir)
+ }
+
+ p, err := profile.Merge(profiles)
+ if err != nil {
+ return "", fmt.Errorf("error merging profiles: %w", err)
+ }
+
+ out := filepath.Join(dir, "merged.cpu")
+ f, err := os.Create(out)
+ defer f.Close()
+
+ if err := p.Write(f); err != nil {
+ return "", fmt.Errorf("error writing merged profile: %w", err)
+ }
+
+ return out, nil
+}
+
 func canonicalizePath(path, base string) string {
  if filepath.IsAbs(path) {
  return path

sweet/common/config.go

@@ -43,10 +43,11 @@ type ConfigFile struct {
 }
 
 type Config struct {
- Name string `toml:"name"`
- GoRoot string `toml:"goroot"`
- BuildEnv ConfigEnv `toml:"envbuild"`
- ExecEnv ConfigEnv `toml:"envexec"`
+ Name string `toml:"name"`
+ GoRoot string `toml:"goroot"`
+ BuildEnv ConfigEnv `toml:"envbuild"`
+ ExecEnv ConfigEnv `toml:"envexec"`
+ PGOFiles map[string]string `toml:"pgofiles"`
 }
 
 func (c *Config) GoTool() *Go {
@@ -58,6 +59,14 @@ func (c *Config) GoTool() *Go {
  }
 }
 
+// Copy returns a deep copy of Config.
+func (c *Config) Copy() *Config {
+ // Currently, all fields in Config are immutable, so a simply copy is
+ // sufficient.
+ cc := *c
+ return &cc
+}
+
 func ConfigFileMarshalTOML(c *ConfigFile) ([]byte, error) {
  // Unfortunately because the github.com/BurntSushi/toml
  // package at v1.0.0 doesn't correctly support Marshaler
@@ -67,10 +76,11 @@ func ConfigFileMarshalTOML(c *ConfigFile) ([]byte, error) {
  // on Config and use dummy types that have a straightforward
  // mapping that *does* work.
  type config struct {
- Name string `toml:"name"`
- GoRoot string `toml:"goroot"`
- BuildEnv []string `toml:"envbuild"`
- ExecEnv []string `toml:"envexec"`
+ Name string `toml:"name"`
+ GoRoot string `toml:"goroot"`
+ BuildEnv []string `toml:"envbuild"`
+ ExecEnv []string `toml:"envexec"`
+ PGOFiles map[string]string `toml:"pgofiles"`
  }
  type configFile struct {
  Configs []*config `toml:"config"`
@@ -82,6 +92,7 @@ func ConfigFileMarshalTOML(c *ConfigFile) ([]byte, error) {
  cfg.GoRoot = c.GoRoot
  cfg.BuildEnv = c.BuildEnv.Collapse()
  cfg.ExecEnv = c.ExecEnv.Collapse()
+ cfg.PGOFiles = c.PGOFiles
 
  cfgs.Configs = append(cfgs.Configs, &cfg)
  }