diff --git a/cc/action_names.bzl b/cc/action_names.bzl index 7689a4c1a..2b734ea2d 100644 --- a/cc/action_names.bzl +++ b/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, diff --git a/cc/private/rules_impl/cc_toolchain_provider_helper.bzl b/cc/private/rules_impl/cc_toolchain_provider_helper.bzl index 33a088cb9..42ba78a90 100644 --- a/cc/private/rules_impl/cc_toolchain_provider_helper.bzl +++ b/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 diff --git a/cc/private/rules_impl/fdo/fdo_context.bzl b/cc/private/rules_impl/fdo/fdo_context.bzl index e260ae93d..d593114bb 100644 --- a/cc/private/rules_impl/fdo/fdo_context.bzl +++ b/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}", ) diff --git a/cc/toolchains/actions/BUILD b/cc/toolchains/actions/BUILD index b231568c6..871373670 100644 --- a/cc/toolchains/actions/BUILD +++ b/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", diff --git a/tests/cc/common/BUILD b/tests/cc/common/BUILD index 7141680a9..832265c51 100644 --- a/tests/cc/common/BUILD +++ b/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") diff --git a/tests/cc/common/cc_fdo_env_test.bzl b/tests/cc/common/cc_fdo_env_test.bzl new file mode 100644 index 000000000..d68f33f15 --- /dev/null +++ b/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, + ) diff --git a/tests/cc/testutil/toolchains/cc_toolchain_config.bzl b/tests/cc/testutil/toolchains/cc_toolchain_config.bzl index 90c0eef90..7b8b2355d 100644 --- a/tests/cc/testutil/toolchains/cc_toolchain_config.bzl +++ b/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, diff --git a/tests/cc/testutil/toolchains/features.bzl b/tests/cc/testutil/toolchains/features.bzl index 75625eb35..34562e847 100644 --- a/tests/cc/testutil/toolchains/features.bzl +++ b/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", )