Add pgo support for go 1.20 (#3641)

* Add pgo support for go 1.20

* Pull request feedback from @fmeum

* Add test for building a go_binary with and without a pgoprofile

* Add //go/config:pgoprofile to _common_reset_transition_dict

* Apply suggestions from code review

Co-authored-by: Fabian Meumertzheim <fabian@meumertzhe.im>

* Use `cquery --output=files` to simplify `pgo_test`

---------

Co-authored-by: Fabian Meumertzheim <fabian@meumertzhe.im>

BUILD.bazel

@@ -123,6 +123,7 @@ go_config(
     gotags = "//go/config:tags",
     linkmode = "//go/config:linkmode",
     msan = "//go/config:msan",
+    pgoprofile = "//go/config:pgoprofile",
     pure = "//go/config:pure",
     race = "//go/config:race",
     stamp = select({

docs/go/core/rules.md

@@ -126,7 +126,7 @@ Rules
 <pre>
 go_binary(<a href="#go_binary-name">name</a>, <a href="#go_binary-basename">basename</a>, <a href="#go_binary-cdeps">cdeps</a>, <a href="#go_binary-cgo">cgo</a>, <a href="#go_binary-clinkopts">clinkopts</a>, <a href="#go_binary-copts">copts</a>, <a href="#go_binary-cppopts">cppopts</a>, <a href="#go_binary-cxxopts">cxxopts</a>, <a href="#go_binary-data">data</a>, <a href="#go_binary-deps">deps</a>, <a href="#go_binary-embed">embed</a>,
           <a href="#go_binary-embedsrcs">embedsrcs</a>, <a href="#go_binary-env">env</a>, <a href="#go_binary-gc_goopts">gc_goopts</a>, <a href="#go_binary-gc_linkopts">gc_linkopts</a>, <a href="#go_binary-goarch">goarch</a>, <a href="#go_binary-goos">goos</a>, <a href="#go_binary-gotags">gotags</a>, <a href="#go_binary-importpath">importpath</a>, <a href="#go_binary-linkmode">linkmode</a>, <a href="#go_binary-msan">msan</a>,
-          <a href="#go_binary-out">out</a>, <a href="#go_binary-pure">pure</a>, <a href="#go_binary-race">race</a>, <a href="#go_binary-srcs">srcs</a>, <a href="#go_binary-static">static</a>, <a href="#go_binary-x_defs">x_defs</a>)
+          <a href="#go_binary-out">out</a>, <a href="#go_binary-pgoprofile">pgoprofile</a>, <a href="#go_binary-pure">pure</a>, <a href="#go_binary-race">race</a>, <a href="#go_binary-srcs">srcs</a>, <a href="#go_binary-static">static</a>, <a href="#go_binary-x_defs">x_defs</a>)
 </pre>
 
 This builds an executable from a set of source files,
@@ -168,6 +168,7 @@ This builds an executable from a set of source files,
 | <a id="go_binary-linkmode"></a>linkmode |  Determines how the binary should be built and linked. This accepts some of             the same values as `go build -buildmode` and works the same way.             <br><br>             <ul>             <li>`auto` (default): Controlled by `//go/config:linkmode`, which defaults to `normal`.</li>             <li>`normal`: Builds a normal executable with position-dependent code.</li>             <li>`pie`: Builds a position-independent executable.</li>             <li>`plugin`: Builds a shared library that can be loaded as a Go plugin. Only supported on platforms that support plugins.</li>             <li>`c-shared`: Builds a shared library that can be linked into a C program.</li>             <li>`c-archive`: Builds an archive that can be linked into a C program.</li>             </ul>   | String | optional | "auto" |
 | <a id="go_binary-msan"></a>msan |  Controls whether code is instrumented for memory sanitization. May be one of             <code>on</code>, <code>off</code>, or <code>auto</code>. Not available when cgo is             disabled. In most cases, it's better to control this on the command line with             <code>--@io_bazel_rules_go//go/config:msan</code>. See [mode attributes], specifically             [msan].   | String | optional | "auto" |
 | <a id="go_binary-out"></a>out |  Sets the output filename for the generated executable. When set, <code>go_binary</code>             will write this file without mode-specific directory prefixes, without             linkmode-specific prefixes like "lib", and without platform-specific suffixes             like ".exe". Note that without a mode-specific directory prefix, the             output file (but not its dependencies) will be invalidated in Bazel's cache             when changing configurations.   | String | optional | "" |
+| <a id="go_binary-pgoprofile"></a>pgoprofile |  Provides a pprof file to be used for profile guided optimization when compiling go targets.             A pprof file can also be provided via <code>--@io_bazel_rules_go//go/config:pgoprofile=&lt;label of a pprof file&gt;</code>.             Profile guided optimization is only supported on go 1.20+.             See https://go.dev/doc/pgo for more information.   | <a href="https://bazel.build/concepts/labels">Label</a> | optional | //go/config:empty |
 | <a id="go_binary-pure"></a>pure |  Controls whether cgo source code and dependencies are compiled and linked,             similar to setting <code>CGO_ENABLED</code>. May be one of <code>on</code>, <code>off</code>,             or <code>auto</code>. If <code>auto</code>, pure mode is enabled when no C/C++             toolchain is configured or when cross-compiling. It's usually better to             control this on the command line with             <code>--@io_bazel_rules_go//go/config:pure</code>. See [mode attributes], specifically             [pure].   | String | optional | "auto" |
 | <a id="go_binary-race"></a>race |  Controls whether code is instrumented for race detection. May be one of             <code>on</code>, <code>off</code>, or <code>auto</code>. Not available when cgo is             disabled. In most cases, it's better to control this on the command line with             <code>--@io_bazel_rules_go//go/config:race</code>. See [mode attributes], specifically             [race].   | String | optional | "auto" |
 | <a id="go_binary-srcs"></a>srcs |  The list of Go source files that are compiled to create the package.             Only <code>.go</code> and <code>.s</code> files are permitted, unless the <code>cgo</code>             attribute is set, in which case,             <code>.c .cc .cpp .cxx .h .hh .hpp .hxx .inc .m .mm</code>             files are also permitted. Files may be filtered at build time             using Go [build constraints].   | <a href="https://bazel.build/concepts/labels">List of labels</a> | optional | [] |

go/config/BUILD.bazel

@@ -79,3 +79,14 @@ string_list_flag(
     build_setting_default = [],
     visibility = ["//visibility:public"],
 )
+
+label_flag(
+    name = "pgoprofile",
+    build_setting_default = ":empty",
+    visibility = ["//visibility:public"],
+)
+
+filegroup(
+    name = "empty",
+    visibility = ["//visibility:public"],
+)

go/private/actions/compilepkg.bzl

@@ -151,6 +151,10 @@ def emit_compilepkg(
         if clinkopts:
             args.add("-ldflags", quote_opts(clinkopts))
 
+    if go.mode.pgoprofile:
+        args.add("-pgoprofile", go.mode.pgoprofile)
+        inputs.append(go.mode.pgoprofile)
+
     go.actions.run(
         inputs = inputs,
         outputs = outputs,

go/private/actions/stdlib.bzl

@@ -136,6 +136,11 @@ def _build_stdlib(go):
               go.sdk.tools +
               [go.sdk.go, go.sdk.package_list, go.sdk.root_file] +
               go.crosstool)
+
+    if go.mode.pgoprofile:
+        args.add("-pgoprofile", go.mode.pgoprofile)
+        inputs.append(go.mode.pgoprofile)
+
     outputs = [pkg]
     go.actions.run(
         inputs = inputs,

go/private/context.bzl

@@ -269,6 +269,7 @@ def _library_to_source(go, attr, library, coverage_instrumented):
         "cgo_deps": [],
         "cgo_exports": [],
         "cc_info": None,
+        "pgoprofile": getattr(attr, "pgoprofile", None),
     }
     if coverage_instrumented:
         source["cover"] = attr_srcs
@@ -530,6 +531,7 @@ def go_context(ctx, attr = None):
         stamp = mode.stamp,
         label = ctx.label,
         cover_format = mode.cover_format,
+        pgoprofile = mode.pgoprofile,
         # Action generators
         archive = toolchain.actions.archive,
         binary = toolchain.actions.binary,
@@ -836,6 +838,7 @@ def _go_config_impl(ctx):
         cover_format = ctx.attr.cover_format[BuildSettingInfo].value,
         gc_goopts = ctx.attr.gc_goopts[BuildSettingInfo].value,
         amd64 = ctx.attr.amd64,
+        pgoprofile = ctx.attr.pgoprofile,
     )]
 
 go_config = rule(
@@ -884,6 +887,10 @@ go_config = rule(
             providers = [BuildSettingInfo],
         ),
         "amd64": attr.string(),
+        "pgoprofile": attr.label(
+            mandatory = True,
+            allow_files = True,
+        ),
     },
     provides = [GoConfigInfo],
     doc = """Collects information about build settings in the current

go/private/mode.bzl

@@ -93,6 +93,12 @@ def get_mode(ctx, go_toolchain, cgo_context_info, go_config_info):
     goos = go_toolchain.default_goos if getattr(ctx.attr, "goos", "auto") == "auto" else ctx.attr.goos
     goarch = go_toolchain.default_goarch if getattr(ctx.attr, "goarch", "auto") == "auto" else ctx.attr.goarch
     gc_goopts = go_config_info.gc_goopts if go_config_info else []
+    pgoprofile = None
+    if go_config_info:
+        if len(go_config_info.pgoprofile.files.to_list()) > 2:
+            fail("providing more than one pprof file to pgoprofile is not supported")
+        elif len(go_config_info.pgoprofile.files.to_list()) == 1:
+            pgoprofile = go_config_info.pgoprofile.files.to_list()[0]
 
     # TODO(jayconrod): check for more invalid and contradictory settings.
     if pure and race:
@@ -130,6 +136,7 @@ def get_mode(ctx, go_toolchain, cgo_context_info, go_config_info):
         cover_format = cover_format,
         amd64 = amd64,
         gc_goopts = gc_goopts,
+        pgoprofile = pgoprofile,
     )
 
 def installsuffix(mode):

go/private/rules/binary.bzl

@@ -400,6 +400,15 @@ _go_binary_kwargs = {
             </ul>
             """,
         ),
+        "pgoprofile": attr.label(
+            allow_files = True,
+            doc = """Provides a pprof file to be used for profile guided optimization when compiling go targets.
+            A pprof file can also be provided via `--@io_bazel_rules_go//go/config:pgoprofile=<label of a pprof file>`.
+            Profile guided optimization is only supported on go 1.20+.
+            See https://go.dev/doc/pgo for more information.
+            """,
+            default = "//go/config:empty",
+        ),
         "_go_context_data": attr.label(default = "//:go_context_data", cfg = go_transition),
         "_allowlist_function_transition": attr.label(
             default = "@bazel_tools//tools/allowlists/function_transition_allowlist",

go/private/rules/transition.bzl

@@ -43,6 +43,7 @@ TRANSITIONED_GO_SETTING_KEYS = [
     "//go/config:pure",
     "//go/config:linkmode",
     "//go/config:tags",
+    "//go/config:pgoprofile",
 ]
 
 def _deduped_and_sorted(strs):
@@ -117,6 +118,10 @@ def _go_transition_impl(settings, attr):
             fail("linkmode: invalid mode {}; want one of {}".format(linkmode, ", ".join(LINKMODES)))
         settings["//go/config:linkmode"] = linkmode
 
+    pgoprofile = getattr(attr, "pgoprofile", "auto")
+    if pgoprofile != "auto":
+        settings["//go/config:pgoprofile"] = pgoprofile
+
     for key, original_key in _SETTING_KEY_TO_ORIGINAL_SETTING_KEY.items():
         old_value = original_settings[key]
         value = settings[key]
@@ -132,6 +137,9 @@ def _go_transition_impl(settings, attr):
             # original setting wasn't set explicitly (empty string) or was set
             # explicitly to its default  (always a non-empty string with JSON
             # encoding, e.g. "\"\"" or "[]").
+            if type(old_value) == "Label":
+                # Label is not JSON serializable, so we need to convert it to a string.
+                old_value = str(old_value)
             settings[original_key] = json.encode(old_value)
         else:
             settings[original_key] = ""
@@ -177,6 +185,7 @@ _common_reset_transition_dict = dict({
     "//go/config:debug": False,
     "//go/config:linkmode": LINKMODE_NORMAL,
     "//go/config:tags": [],
+    "//go/config:pgoprofile": Label("//go/config:empty"),
 }, **{setting: "" for setting in _SETTING_KEY_TO_ORIGINAL_SETTING_KEY.values()})
 
 _reset_transition_dict = dict(_common_reset_transition_dict, **{
@@ -191,6 +200,7 @@ _stdlib_keep_keys = sorted([
     "//go/config:pure",
     "//go/config:linkmode",
     "//go/config:tags",
+    "//go/config:pgoprofile",
 ])
 
 def _go_tool_transition_impl(settings, _attr):

go/tools/builders/compilepkg.go

@@ -56,6 +56,7 @@ func compilePkg(args []string) error {
 	var testFilter string
 	var gcFlags, asmFlags, cppFlags, cFlags, cxxFlags, objcFlags, objcxxFlags, ldFlags quoteMultiFlag
 	var coverFormat string
+	var pgoprofile string
 	fs.Var(&unfilteredSrcs, "src", ".go, .c, .cc, .m, .mm, .s, or .S file to be filtered and compiled")
 	fs.Var(&coverSrcs, "cover", ".go file that should be instrumented for coverage (must also be a -src)")
 	fs.Var(&embedSrcs, "embedsrc", "file that may be compiled into the package with a //go:embed directive")
@@ -81,6 +82,7 @@ func compilePkg(args []string) error {
 	fs.StringVar(&testFilter, "testfilter", "off", "Controls test package filtering")
 	fs.StringVar(&coverFormat, "cover_format", "", "Emit source file paths in coverage instrumentation suitable for the specified coverage format")
 	fs.Var(&recompileInternalDeps, "recompile_internal_deps", "The import path of the direct dependencies that needs to be recompiled.")
+	fs.StringVar(&pgoprofile, "pgoprofile", "", "The pprof profile to consider for profile guided optimization.")
 	if err := fs.Parse(args); err != nil {
 		return err
 	}
@@ -99,6 +101,9 @@ func compilePkg(args []string) error {
 	for i := range embedSrcs {
 		embedSrcs[i] = abs(embedSrcs[i])
 	}
+	if pgoprofile != "" {
+		pgoprofile = abs(pgoprofile)
+	}
 
 	// Filter sources.
 	srcs, err := filterAndSplitFiles(unfilteredSrcs)
@@ -158,7 +163,8 @@ func compilePkg(args []string) error {
 		outFactsPath,
 		cgoExportHPath,
 		coverFormat,
-		recompileInternalDeps)
+		recompileInternalDeps,
+		pgoprofile)
 }
 
 func compileArchive(
@@ -189,6 +195,7 @@ func compileArchive(
 	cgoExportHPath string,
 	coverFormat string,
 	recompileInternalDeps []string,
+	pgoprofile string,
 ) error {
 	workDir, cleanup, err := goenv.workDir()
 	if err != nil {
@@ -467,7 +474,7 @@ func compileArchive(
 	}
 
 	// Compile the filtered .go files.
-	if err := compileGo(goenv, goSrcs, packagePath, importcfgPath, embedcfgPath, asmHdrPath, symabisPath, gcFlags, outPath); err != nil {
+	if err := compileGo(goenv, goSrcs, packagePath, importcfgPath, embedcfgPath, asmHdrPath, symabisPath, gcFlags, pgoprofile, outPath); err != nil {
 		return err
 	}
 
@@ -537,7 +544,7 @@ func compileArchive(
 	return appendFiles(goenv, outXPath, []string{pkgDefPath})
 }
 
-func compileGo(goenv *env, srcs []string, packagePath, importcfgPath, embedcfgPath, asmHdrPath, symabisPath string, gcFlags []string, outPath string) error {
+func compileGo(goenv *env, srcs []string, packagePath, importcfgPath, embedcfgPath, asmHdrPath, symabisPath string, gcFlags []string, pgoprofile string, outPath string) error {
 	args := goenv.goTool("compile")
 	args = append(args, "-p", packagePath, "-importcfg", importcfgPath, "-pack")
 	if embedcfgPath != "" {
@@ -549,6 +556,9 @@ func compileGo(goenv *env, srcs []string, packagePath, importcfgPath, embedcfgPa
 	if symabisPath != "" {
 		args = append(args, "-symabis", symabisPath)
 	}
+	if pgoprofile != "" {
+		args = append(args, "-pgoprofile", pgoprofile)
+	}
 	args = append(args, gcFlags...)
 	args = append(args, "-o", outPath)
 	args = append(args, "--")

go/tools/builders/stdlib.go

@@ -33,6 +33,7 @@ func stdlib(args []string) error {
 	race := flags.Bool("race", false, "Build in race mode")
 	shared := flags.Bool("shared", false, "Build in shared mode")
 	dynlink := flags.Bool("dynlink", false, "Build in dynlink mode")
+	pgoprofile := flags.String("pgoprofile", "", "Build with pgo using the given pprof file")
 	var packages multiFlag
 	flags.Var(&packages, "package", "Packages to build")
 	var gcflags quoteMultiFlag
@@ -130,6 +131,9 @@ You may need to use the flags --cpu=x64_windows --compiler=mingw-gcc.`)
 	if *race {
 		installArgs = append(installArgs, "-race")
 	}
+	if *pgoprofile != "" {
+		installArgs = append(installArgs, "-pgo", abs(*pgoprofile))
+	}
 	if *shared {
 		gcflags = append(gcflags, "-shared")
 		ldflags = append(ldflags, "-shared")

tests/core/go_binary/BUILD.bazel

@@ -216,3 +216,17 @@ go_bazel_test(
     name = "non_executable_test",
     srcs = ["non_executable_test.go"],
 )
+
+exports_files(["pgo.pprof"])
+
+go_binary(
+    name = "pgo",
+    srcs = ["pgo.go"],
+    pgoprofile = "pgo.pprof",
+)
+
+go_bazel_test(
+    name = "pgo_test",
+    srcs = ["pgo_test.go"],
+    embedsrcs = ["pgo.pprof"],
+)

tests/core/go_binary/pgo.go

@@ -0,0 +1,42 @@
+package main
+
+import (
+	"bytes"
+	"io"
+	"log"
+	"net/http"
+	_ "net/http/pprof"
+)
+
+func render(w http.ResponseWriter, r *http.Request) {
+	if r.Method != "POST" {
+		http.Error(w, "Only POST allowed", http.StatusMethodNotAllowed)
+		return
+	}
+
+	src, err := io.ReadAll(r.Body)
+	if err != nil {
+		log.Printf("error reading body: %v", err)
+		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
+		return
+	}
+
+	var buf bytes.Buffer
+	if _, err := buf.Write(src); err != nil {
+		log.Printf("error writing response: %v", err)
+		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
+		return
+	}
+
+	if _, err := io.Copy(w, &buf); err != nil {
+		log.Printf("error writing response: %v", err)
+		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
+		return
+	}
+}
+
+func main() {
+	http.HandleFunc("/render", render)
+	log.Printf("Serving on port 8080...")
+	log.Fatal(http.ListenAndServe(":8080", nil))
+}