Add llvm-profdata action, support env

cc/action_names.bzl

@@ -52,6 +52,9 @@ PREPROCESS_ASSEMBLE_ACTION_NAME = "preprocess-assemble"
 # Name of the coverage action.
 LLVM_COV = "llvm-cov"
 
+# Name of the profile data merging action.
+LLVM_PROFDATA = "llvm-profdata"
+
 # Name of the action producing ThinLto index.
 LTO_INDEXING_ACTION_NAME = "lto-indexing"
 
@@ -126,6 +129,7 @@ ACTION_NAMES = struct(
     assemble = ASSEMBLE_ACTION_NAME,
     preprocess_assemble = PREPROCESS_ASSEMBLE_ACTION_NAME,
     llvm_cov = LLVM_COV,
+    llvm_profdata = LLVM_PROFDATA,
     lto_indexing = LTO_INDEXING_ACTION_NAME,
     lto_backend = LTO_BACKEND_ACTION_NAME,
     cpp_header_analysis = CPP_HEADER_ANALYSIS_ACTION_NAME,

cc/private/rules_impl/cc_toolchain_provider_helper.bzl

@@ -184,12 +184,16 @@ def get_cc_toolchain_provider(ctx, attributes):
     )
     tool_paths = _compute_tool_paths(toolchain_config_info, tools_directory)
     toolchain_features = cc_common.cc_toolchain_features(toolchain_config_info = toolchain_config_info, tools_directory = tools_directory)
+    feature_configuration = toolchain_features.configure_features(
+        requested_features = toolchain_features.default_features_and_action_configs(),
+    )
     fdo_context = create_fdo_context(
         llvm_profdata = tool_paths.get("llvm-profdata"),
         all_files = attributes.all_files,
         zipper = attributes.zipper,
         cc_toolchain_config_info = toolchain_config_info,
         coverage_enabled = ctx.configuration.coverage_enabled,
+        feature_configuration = feature_configuration,
     )
     if fdo_context == None:
         return None

cc/private/rules_impl/fdo/fdo_context.bzl

@@ -14,6 +14,7 @@
 """FDO context describes how C++ FDO compilation should be done."""
 
 load("@bazel_skylib//lib:paths.bzl", "paths")
+load("//cc:action_names.bzl", "ACTION_NAMES")
 load("//cc/common:cc_common.bzl", "cc_common")
 load("//cc/private/rules_impl/fdo:fdo_prefetch_hints.bzl", "FdoPrefetchHintsInfo")
 load("//cc/private/rules_impl/fdo:fdo_profile.bzl", "FdoProfileInfo")
@@ -28,6 +29,7 @@ def _create_fdo_context(
         zipper,
         cc_toolchain_config_info,
         coverage_enabled,
+        feature_configuration,
         _fdo_prefetch_hints,
         _propeller_optimize,
         _memprof_profile,
@@ -48,6 +50,7 @@ def _create_fdo_context(
       zipper: (File) zip tool, used to unpact the profiles
       cc_toolchain_config_info: (CcToolchainConfigInfo) Used to check CPU value, should be removed
       coverage_enabled: (bool) Is code coverage enabled
+      feature_configuration: (FeatureConfiguration) Used for llvm-profdata action
       _fdo_prefetch_hints: (Target) Pointed to by --fdo_prefetch_hints
       _propeller_optimize: (Target) Pointed to by --propeller_optimize
       _memprof_profile: (Target) Pointed to by --memprof_profile
@@ -70,6 +73,21 @@ def _create_fdo_context(
     if cpp_config.compilation_mode() != "opt":
         return struct()
 
+    llvm_profdata_env = {}
+    if cc_common.action_is_enabled(
+        feature_configuration = feature_configuration,
+        action_name = ACTION_NAMES.llvm_profdata,
+    ):
+        llvm_profdata = cc_common.get_tool_for_action(
+            feature_configuration = feature_configuration,
+            action_name = ACTION_NAMES.llvm_profdata,
+        )
+        llvm_profdata_env = cc_common.get_environment_variables(
+            feature_configuration = feature_configuration,
+            action_name = ACTION_NAMES.llvm_profdata,
+            variables = cc_common.empty_variables(),
+        )
+
     # Propeller optimize cc and ld profiles
     cc_profile = _symlink_to(
         ctx,
@@ -186,6 +204,7 @@ def _create_fdo_context(
                 all_files,
                 zipper,
                 cc_toolchain_config_info,
+                llvm_profdata_env,
             )
         elif branch_fdo_mode in ["auto_fdo", "xbinary_fdo"]:
             profile_artifact = _symlink_input(
@@ -203,6 +222,7 @@ def _create_fdo_context(
                 all_files,
                 zipper,
                 cc_toolchain_config_info,
+                llvm_profdata_env,
             )
             cs_profile_artifact = _convert_llvm_raw_profile_to_indexed(
                 ctx,
@@ -212,6 +232,7 @@ def _create_fdo_context(
                 all_files,
                 zipper,
                 cc_toolchain_config_info,
+                llvm_profdata_env,
             )
             profile_artifact = _merge_llvm_profiles(
                 ctx,
@@ -221,6 +242,7 @@ def _create_fdo_context(
                 non_cs_profile_artifact,
                 cs_profile_artifact,
                 "MergedCS.profdata",
+                llvm_profdata_env,
             )
 
         branch_fdo_profile = struct(
@@ -254,7 +276,8 @@ def _convert_llvm_raw_profile_to_indexed(
         llvm_profdata,
         all_files,
         zipper,
-        cc_toolchain_config_info):
+        cc_toolchain_config_info,
+        env):
     """This function checks the input profile format and converts it to the indexed format (.profdata) if necessary."""
     basename = _basename(fdo_inputs)
     if basename.endswith(".profdata"):
@@ -310,6 +333,7 @@ def _convert_llvm_raw_profile_to_indexed(
         arguments = [ctx.actions.args().add("merge").add("-o").add(profile_artifact).add(raw_profile_artifact)],
         inputs = [raw_profile_artifact],
         outputs = [profile_artifact],
+        env = env,
         use_default_shell_env = True,
         progress_message = "LLVMProfDataAction: Generating %{output}",
     )
@@ -323,7 +347,8 @@ def _merge_llvm_profiles(
         all_files,
         profile1,
         profile2,
-        merged_output_name):
+        merged_output_name,
+        env):
     """This function merges profile1 and profile2 and generates merged_output."""
     profile_artifact = ctx.actions.declare_file(name_prefix + "/" + ctx.label.name + "/" + merged_output_name)
 
@@ -335,6 +360,7 @@ def _merge_llvm_profiles(
         arguments = [ctx.actions.args().add("merge").add("-o").add(profile_artifact).add(profile1).add(profile2)],
         inputs = [profile1, profile2],
         outputs = [profile_artifact],
+        env = env,
         use_default_shell_env = True,
         progress_message = "LLVMProfDataAction: Generating %{output}",
     )

cc/toolchains/actions/BUILD

@@ -88,6 +88,11 @@ cc_action_type(
     action_name = ACTION_NAMES.llvm_cov,
 )
 
+cc_action_type(
+    name = "llvm_profdata",
+    action_name = ACTION_NAMES.llvm_profdata,
+)
+
 cc_action_type(
     name = "lto_indexing",
     action_name = ACTION_NAMES.lto_indexing,
@@ -324,6 +329,7 @@ cc_action_type_set(
         ":assemble",
         ":preprocess_assemble",
         ":llvm_cov",
+        ":llvm_profdata",
         ":lto_indexing",
         ":lto_backend",
         ":lto_index_for_executable",

tests/cc/common/BUILD

@@ -1,6 +1,9 @@
 load(":cc_binary_configured_target_tests.bzl", "cc_binary_configured_target_tests")
 load(":cc_common_test.bzl", "cc_common_tests")
+load(":cc_fdo_env_test.bzl", "cc_fdo_env_tests")
 
 cc_binary_configured_target_tests(name = "cc_binary_configured_target_tests")
 
 cc_common_tests(name = "cc_common_tests")
+
+cc_fdo_env_tests(name = "cc_fdo_env_tests")

tests/cc/common/cc_fdo_env_test.bzl

@@ -0,0 +1,119 @@
+"""Tests for FDO action environment variables."""
+
+load("@bazel_features//private:util.bzl", _bazel_version_ge = "ge")
+load("@rules_testing//lib:analysis_test.bzl", "analysis_test", "test_suite")
+load("@rules_testing//lib:util.bzl", "TestingAspectInfo")
+load("//cc:action_names.bzl", "ACTION_NAMES")
+load("//cc/toolchains:fdo_profile.bzl", "fdo_profile")
+load("//tests/cc/testutil/toolchains:features.bzl", "FEATURE_NAMES")
+
+def _test_fdo_profdata_env(name, **kwargs):
+    native.genrule(
+        name = name + "_profraw",
+        outs = [name + "/profile.profraw"],
+        cmd = "touch $@",
+    )
+
+    fdo_profile(
+        name = name + "_profile",
+        profile = name + "_profraw",
+    )
+
+    analysis_test(
+        name = name,
+        impl = _test_fdo_profdata_env_impl,
+        target = "//tests/cc/testutil/toolchains:cc-compiler-k8-compiler",
+        config_settings = _fdo_config_settings(name),
+        **kwargs
+    )
+
+def _fdo_config_settings(name):
+    return {
+        str(Label("//tests/cc/testutil/toolchains:with_features")): [
+            FEATURE_NAMES.fdo_optimize,
+            FEATURE_NAMES.llvm_profdata_env,
+        ],
+        str(Label("//tests/cc/testutil/toolchains:with_action_configs")): [
+            ACTION_NAMES.llvm_profdata,
+        ],
+        "//command_line_option:fdo_optimize": "//tests/cc/common:" + name + "_profile",
+        "//command_line_option:compilation_mode": "opt",
+    }
+
+def _assert_profdata_env(env, target, expected_mnemonics):
+    profdata_actions = []
+    seen_mnemonics = {}
+    for action in target[TestingAspectInfo].actions:
+        if action.mnemonic in expected_mnemonics:
+            profdata_actions.append(action)
+            seen_mnemonics[action.mnemonic] = True
+
+    env.expect.that_collection(seen_mnemonics.keys()).contains_at_least(expected_mnemonics)
+    for action in profdata_actions:
+        env.expect.where(
+            detail = "mnemonic: %s" % action.mnemonic,
+        ).that_dict(action.env).contains_at_least({"LLVM_PROFDATA_ENV_KEY": "LLVM_PROFDATA_ENV_VALUE"})
+
+        env.expect.where(
+            detail = "mnemonic: %s" % action.mnemonic,
+        ).that_dict(action.env).keys().contains("PATH")
+
+def _test_fdo_profdata_env_impl(env, target):
+    _assert_profdata_env(env, target, ["LLVMProfDataAction"])
+
+def _test_csfdo_profdata_merge_env(name, **kwargs):
+    """Tests that LLVMProfDataMergeAction gets env vars from the feature configuration.
+
+    CS-FDO requires both --fdo_optimize and --cs_fdo_profile, triggering a merge action.
+    """
+    native.genrule(
+        name = name + "_profraw",
+        outs = [name + "/profile.profraw"],
+        cmd = "touch $@",
+    )
+
+    fdo_profile(
+        name = name + "_profile",
+        profile = name + "_profraw",
+    )
+
+    config = _fdo_config_settings(name)
+    config["//command_line_option:cs_fdo_profile"] = str(Label("//tests/cc/common:cs_fdo_profile"))
+    analysis_test(
+        name = name,
+        impl = _test_csfdo_profdata_merge_env_impl,
+        target = "//tests/cc/testutil/toolchains:cc-compiler-k8-compiler",
+        config_settings = config,
+        **kwargs
+    )
+
+def _test_csfdo_profdata_merge_env_impl(env, target):
+    _assert_profdata_env(env, target, ["LLVMProfDataMergeAction"])
+
+def cc_fdo_env_tests(name):
+    """Creates tests for FDO actions.
+
+    Args:
+        name: The name of the test suite.
+    """
+    native.genrule(
+        name = "cs_fdo_profraw",
+        outs = ["cs_profile.profraw"],
+        cmd = "touch $@",
+    )
+
+    fdo_profile(
+        name = "cs_fdo_profile",
+        profile = ":cs_fdo_profraw",
+    )
+
+    tests = []
+    if _bazel_version_ge("9.0.0"):
+        tests = [
+            _test_fdo_profdata_env,
+            _test_csfdo_profdata_merge_env,
+        ]
+    test_suite(
+        name = name,
+        tests = tests,
+    )

tests/cc/testutil/toolchains/cc_toolchain_config.bzl

@@ -832,6 +832,22 @@ _simple_link_feature = feature(
     ],
 )
 
+_llvm_profdata_env_feature = feature(
+    name = FEATURE_NAMES.llvm_profdata_env,
+    enabled = True,
+    env_sets = [
+        env_set(
+            actions = [ACTION_NAMES.llvm_profdata],
+            env_entries = [
+                env_entry(
+                    key = "LLVM_PROFDATA_ENV_KEY",
+                    value = "LLVM_PROFDATA_ENV_VALUE",
+                ),
+            ],
+        ),
+    ],
+)
+
 _link_env_feature = feature(
     name = FEATURE_NAMES.link_env,
     env_sets = [
@@ -1341,6 +1357,7 @@ _feature_name_to_feature = {
     FEATURE_NAMES.simple_compile_feature: _simple_compile_feature,
     FEATURE_NAMES.simple_link_feature: _simple_link_feature,
     FEATURE_NAMES.link_env: _link_env_feature,
+    FEATURE_NAMES.llvm_profdata_env: _llvm_profdata_env_feature,
     FEATURE_NAMES.static_linking_mode: _static_linking_mode_feature,
     FEATURE_NAMES.dynamic_linking_mode: _dynamic_linking_mode_feature,
     FEATURE_NAMES.objcopy_embed_flags: _objcopy_embed_flags_feature,

tests/cc/testutil/toolchains/features.bzl

@@ -94,4 +94,5 @@ FEATURE_NAMES = struct(
     no_copts_tokenization = "no_copts_tokenization",
     generate_linkmap = "generate_linkmap",
     shorten_virtual_includes = "shorten_virtual_includes",
+    llvm_profdata_env = "llvm_profdata_env",
 )