From 90886c730bf69421e932d95da13d8cdb30b7203b Mon Sep 17 00:00:00 2001
From: Fabian Meumertzheim <fabian@meumertzhe.im>
Date: Tue, 12 Dec 2023 20:29:26 +0100
Subject: [PATCH] feat: add support for PGO (#205)

The new `profiles` attribute on `native_image` can be used to apply
profile-guided optimization.

Signed-off-by: Fabian Meumertzheim <fabian@meumertzhe.im>
---
 docs/api/defs.md                  | 3 ++-
 graalvm/nativeimage/rules.bzl     | 3 +++
 internal/native_image/builder.bzl | 8 ++++++++
 internal/native_image/common.bzl  | 4 ++++
 4 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/docs/api/defs.md b/docs/api/defs.md
index 92d6e631..01490fec 100755
--- a/docs/api/defs.md
+++ b/docs/api/defs.md
@@ -11,7 +11,7 @@ native_image(<a href="#native_image-name">name</a>, <a href="#native_image-deps"
              <a href="#native_image-jni_configuration">jni_configuration</a>, <a href="#native_image-initialize_at_build_time">initialize_at_build_time</a>, <a href="#native_image-initialize_at_run_time">initialize_at_run_time</a>, <a href="#native_image-native_features">native_features</a>,
              <a href="#native_image-debug">debug</a>, <a href="#native_image-optimization_mode">optimization_mode</a>, <a href="#native_image-shared_library">shared_library</a>, <a href="#native_image-static_zlib">static_zlib</a>, <a href="#native_image-c_compiler_option">c_compiler_option</a>, <a href="#native_image-data">data</a>,
              <a href="#native_image-extra_args">extra_args</a>, <a href="#native_image-allow_fallback">allow_fallback</a>, <a href="#native_image-check_toolchains">check_toolchains</a>, <a href="#native_image-native_image_tool">native_image_tool</a>, <a href="#native_image-native_image_settings">native_image_settings</a>,
-             <a href="#native_image-kwargs">kwargs</a>)
+             <a href="#native_image-profiles">profiles</a>, <a href="#native_image-kwargs">kwargs</a>)
 </pre>
 
 Generates and compiles a GraalVM native image from a Java library target.
@@ -42,6 +42,7 @@ Generates and compiles a GraalVM native image from a Java library target.
 | <a id="native_image-check_toolchains"></a>check_toolchains |  Whether to perform toolchain checks in `native-image`; defaults to `True` on Windows, `False` otherwise.   |  `select({"@bazel_tools//src/conditions:windows": True, "//conditions:default": False})` |
 | <a id="native_image-native_image_tool"></a>native_image_tool |  Specific `native-image` executable target to use.   |  `None` |
 | <a id="native_image-native_image_settings"></a>native_image_settings |  Suite(s) of Native Image build settings to use.   |  `[Label("@rules_graalvm//internal/native_image:defaults")]` |
+| <a id="native_image-profiles"></a>profiles |  Profiles to use for profile-guided optimization (PGO) and obtained from a native image compiled with `--pgo-instrument`.   |  `[]` |
 | <a id="native_image-kwargs"></a>kwargs |  Extra keyword arguments are passed to the underlying `native_image` rule.   |  none |
 
 
diff --git a/graalvm/nativeimage/rules.bzl b/graalvm/nativeimage/rules.bzl
index fdb4a12c..45ba9a7a 100644
--- a/graalvm/nativeimage/rules.bzl
+++ b/graalvm/nativeimage/rules.bzl
@@ -82,6 +82,7 @@ def native_image(
         check_toolchains = _DEFAULT_CHECK_TOOLCHAINS_CONDITION,
         native_image_tool = None,  # uses toolchains by default
         native_image_settings = [_DEFAULT_NATIVE_IMAGE_SETTINGS],
+        profiles = [],
         **kwargs):
     """Generates and compiles a GraalVM native image from a Java library target.
 
@@ -113,6 +114,7 @@ def native_image(
         check_toolchains: Whether to perform toolchain checks in `native-image`; defaults to `True` on Windows, `False` otherwise.
         native_image_tool: Specific `native-image` executable target to use.
         native_image_settings: Suite(s) of Native Image build settings to use.
+        profiles: Profiles to use for profile-guided optimization (PGO) and obtained from a native image compiled with `--pgo-instrument`.
         **kwargs: Extra keyword arguments are passed to the underlying `native_image` rule.
     """
 
@@ -138,5 +140,6 @@ def native_image(
         executable_name = executable_name,
         native_image_tool = native_image_tool,
         native_image_settings = native_image_settings,
+        profiles = profiles,
         **kwargs
     )
diff --git a/internal/native_image/builder.bzl b/internal/native_image/builder.bzl
index 925976c4..20c717ef 100644
--- a/internal/native_image/builder.bzl
+++ b/internal/native_image/builder.bzl
@@ -208,6 +208,14 @@ def assemble_native_build_options(
             direct_inputs,
         )
 
+    if ctx.files.profiles:
+        args.add_joined(
+            ctx.files.profiles,
+            join_with = ",",
+            format_joined = "--pgo=%s",
+        )
+        direct_inputs.extend(ctx.files.profiles)
+
     # add test-related flags, if this is a `testonly` target
     if ctx.attr.testonly:
         _configure_native_test_flags(
diff --git a/internal/native_image/common.bzl b/internal/native_image/common.bzl
index 5c043a8d..7e3429e4 100644
--- a/internal/native_image/common.bzl
+++ b/internal/native_image/common.bzl
@@ -109,6 +109,10 @@ _NATIVE_IMAGE_ATTRS = {
     "executable_name": attr.string(
         mandatory = True,
     ),
+    "profiles": attr.label_list(
+        allow_files = True,
+        mandatory = False,
+    ),
     "_cc_toolchain": attr.label(
         default = Label(_BAZEL_CURRENT_CPP_TOOLCHAIN),
     ),