SERVER-93201: add PGO (#25342)

GitOrigin-RevId: 2312148c761d206dcc0404824579ee8586225f77

SConstruct

@@ -463,6 +463,18 @@ add_option(
     nargs=0,
 )
 
+add_option(
+    "pgo-profile",
+    help="compile with pgo profiling",
+    nargs=0,
+)
+
+add_option(
+    "pgo",
+    help="compile with pgo. Assumes profile file default.profdata at root of repository",
+    nargs=0,
+)
+
 add_option(
     "enable-http-client",
     choices=["auto", "on", "off"],
@@ -4829,16 +4841,43 @@ def doConfigure(myenv):
                 myenv.ConfError("Failed to enable thin LTO")
 
         if linker_ld != "gold" and not env.TargetOSIs("darwin", "macOS"):
-            myenv.AppendUnique(
-                CCFLAGS=["-ffunction-sections"],
-                LINKFLAGS=[
-                    "-Wl,--symbol-ordering-file=symbols.orderfile",
-                    "-Wl,--no-warn-symbol-ordering",
-                ],
-            )
+            if has_option("pgo"):
+                print("WARNING: skipping symbol ordering as pgo is enabled")
+            else:
+                myenv.AppendUnique(
+                    CCFLAGS=["-ffunction-sections"],
+                    LINKFLAGS=[
+                        "-Wl,--symbol-ordering-file=symbols.orderfile",
+                        "-Wl,--no-warn-symbol-ordering",
+                    ],
+                )
         else:
             print("WARNING: lld linker is required to sort symbols")
 
+        if has_option("pgo-profile"):
+            if (
+                not myenv.ToolchainIs("clang")
+                or not myenv.TargetOSIs("linux")
+                or linker_ld == "gold"
+            ):
+                myenv.FatalError("Error: pgo only works on linux with clang + lld")
+            myenv.AppendUnique(
+                CCFLAGS=["-fprofile-instr-generate"],
+                LINKFLAGS=["-fprofile-instr-generate"],
+            )
+
+        if has_option("pgo"):
+            if (
+                not myenv.ToolchainIs("clang")
+                or not myenv.TargetOSIs("linux")
+                or linker_ld == "gold"
+            ):
+                myenv.FatalError("Error: pgo only works on linux with clang + lld")
+            myenv.AppendUnique(
+                _NON_CONF_CCFLAGS_GEN=["-fprofile-use=./default.profdata"],
+            )
+            myenv["CCFLAGS_WERROR"].remove("-Werror")
+
         # As far as we know these flags only apply on posix-y systems,
         # and not on Darwin.
         if env.TargetOSIs("posix") and not env.TargetOSIs("darwin"):

bazel/config/BUILD.bazel

@@ -20,6 +20,7 @@ load(
     "lsan",
     "msan",
     "opt",
+    "pgo_profile",
     "release",
     "separate_debug",
     "shared_archive",
@@ -1778,3 +1779,19 @@ config_setting(
         "//bazel/config:gcov": "True",
     },
 )
+
+# --------------------------------------
+# pgo_profile options
+# --------------------------------------
+
+pgo_profile(
+    name = "pgo_profile",
+    build_setting_default = False,
+)
+
+config_setting(
+    name = "pgo_profile_enabled",
+    flag_values = {
+        "//bazel/config:pgo_profile": "True",
+    },
+)

bazel/config/configs.bzl

@@ -527,3 +527,17 @@ gcov = rule(
     implementation = lambda ctx: gcov_provider(enabled = ctx.build_setting_value),
     build_setting = config.bool(flag = True),
 )
+
+# =========
+# pgo_profile
+# =========
+
+pgo_profile_provider = provider(
+    doc = "Choose if pgo profiling should be generated",
+    fields = ["enabled"],
+)
+
+pgo_profile = rule(
+    implementation = lambda ctx: pgo_profile_provider(enabled = ctx.build_setting_value),
+    build_setting = config.bool(flag = True),
+)

etc/evergreen_yml_components/definitions.yml

@@ -2647,6 +2647,15 @@ functions:
         args:
           - "./src/evergreen/sasl_windows_cyrussasl_teardown.sh"
 
+  "fetch pgo profile":
+    command: subprocess.exec
+    params:
+      binary: bash
+      args:
+        - "./src/evergreen/functions/download_pgo_profile.sh"
+      env:
+        PGO_PROFILE_URL: ${pgo_profile_url}
+
   "process code coverage data": &process_code_coverage_data
     command: subprocess.exec
     params:

etc/evergreen_yml_components/tasks/compile_tasks.yml

@@ -38,6 +38,7 @@ variables:
       - func: "set task expansion macros"
       - func: "f_expansions_write"
       - func: "get engflow creds"
+      - func: "fetch pgo profile"
     teardown_task:
       - func: "f_expansions_write"
       - func: "attach scons logs"

etc/system_perf.yml

@@ -171,6 +171,20 @@ functions:
         permissions: public-read
         content_type: application/x-gzip
         display_name: DSI Artifacts - Execution ${execution}
+    - command: shell.exec
+      params:
+        script: |
+          cp ./build/Artifacts/DSIArtifacts/WorkloadOutput/reports/ycsb_100read/mongod.0/pre_post_task_output/default.profdata.tgz ./default.profdata.tgz || true
+    - command: s3.put
+      params:
+        aws_key: ${aws_key}
+        aws_secret: ${aws_secret}
+        local_file: ./default.profdata.tgz
+        remote_file: ${project_dir}/${build_variant}/${revision}/${task_id}/${version_id}/logs/default.profdata.tgz
+        bucket: mciuploads
+        permissions: public-read
+        content_type: application/x-gzip
+        display_name: PGO profile data - Execution ${execution}
     - command: s3.put
       params:
         aws_key: ${aws_key}

evergreen/functions/download_pgo_profile.sh

@@ -0,0 +1,14 @@
+DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" > /dev/null 2>&1 && pwd)"
+. "$DIR/../prelude.sh"
+
+set +o errexit
+
+cd src
+
+if [ -z "${PGO_PROFILE_URL:-}" ]; then
+  echo "No pgo profile url specified" >&2
+  exit 0
+fi
+
+wget $PGO_PROFILE_URL
+tar -xvf default.profdata.tgz

site_scons/site_tools/integrate_bazel.py

@@ -769,6 +769,7 @@ def generate(env: SCons.Environment.Environment) -> None:
         f'--//bazel/config:visibility_support={env.GetOption("visibility-support")}',
         f'--//bazel/config:disable_warnings_as_errors={env.GetOption("disable-warnings-as-errors") == "source"}',
         f'--//bazel/config:gcov={env.GetOption("gcov") is not None}',
+        f'--//bazel/config:pgo_profile={env.GetOption("pgo-profile") is not None}',
         f"--platforms=//bazel/platforms:{distro_or_os}_{normalized_arch}_{env.ToolchainName()}",
         f"--host_platform=//bazel/platforms:{distro_or_os}_{normalized_arch}_{env.ToolchainName()}",
         f'--//bazel/config:ssl={"True" if env.GetOption("ssl") == "on" else "False"}',

src/mongo/util/BUILD.bazel

@@ -48,5 +48,8 @@ mongo_cc_library(
     local_defines = select({
         "//bazel/config:gcov_enabled": ["MONGO_GCOV"],
         "//conditions:default": [],
+    }) + select({
+        "//bazel/config:pgo_profile_enabled": ["MONGO_PGO_PROFILE"],
+        "//conditions:default": [],
     }),
 )

src/mongo/util/quick_exit.cpp

@@ -67,6 +67,11 @@
 #include <xray/xray_log_interface.h>
 #endif
 
+#ifdef MONGO_PGO_PROFILE
+extern "C" void __llvm_profile_dump();
+extern "C" void __llvm_profile_reset_counters();
+#endif
+
 #ifdef MONGO_GCOV
 extern "C" void __gcov_flush();
 extern "C" void __gcov_dump();
@@ -94,6 +99,11 @@ void quickExitWithoutLogging(ExitCode code) {
 #endif
 #endif
 
+#ifdef MONGO_PGO_PROFILE
+    __llvm_profile_dump();
+    __llvm_profile_reset_counters();
+#endif
+
 #if __has_feature(xray_instrument)
     // Stop XRay and flush the xray basic log
     /* ignore */ __xray_log_finalize();