From fae114ca2f4da560b6c85da47227d6b4769be8d6 Mon Sep 17 00:00:00 2001 From: Parker Timmerman Date: Mon, 16 Dec 2024 11:54:53 -0500 Subject: [PATCH] feat: Add `//rust/settings:lto` (#3104) Fixes https://github.com/bazelbuild/rules_rust/issues/3045 This PR adds a new build setting `//rust/settings:lto=(off|thin|fat)` which changes how we specify the following flags: * [`lto`](https://doc.rust-lang.org/rustc/codegen-options/index.html#lto) * [`embed-bitcode`](https://doc.rust-lang.org/rustc/codegen-options/index.html#embed-bitcode) * [`linker-plugin-lto`](https://doc.rust-lang.org/rustc/codegen-options/index.html#linker-plugin-lto) The way we invoke the flags was based on how Cargo does it today ([code](https://github.com/rust-lang/cargo/blob/769f622e12db0001431d8ae36d1093fb8727c5d9/src/cargo/core/compiler/lto.rs#L4)) and based on suggestions from the [Rust docs](https://doc.rust-lang.org/rustc/codegen-options/index.html#embed-bitcode). When LTO is not enabled, we will specify `-Cembed-bitcode=no` which tells `rustc` to skip embedding LLVM bitcode and should speed up builds. Similarly when LTO is enabled we specify `-Clinker-plugin-lto` which will cause `rustc` to skip generating objects files entirely, and instead replace them with LLVM bitcode*. *only when building an `rlib`, when building other crate types we continue generating object files. I added unit tests to make sure we pass the flags correctly, as well as some docs describing the new setting. Please let me know if I should add more! --- rust/private/lto.bzl | 114 ++++++++++++++++++++++++++++ rust/private/rustc.bzl | 14 ++++ rust/settings/BUILD.bazel | 7 ++ rust/toolchain.bzl | 7 ++ test/unit/lto/BUILD.bazel | 5 ++ test/unit/lto/lto_test_suite.bzl | 123 +++++++++++++++++++++++++++++++ 6 files changed, 270 insertions(+) create mode 100644 rust/private/lto.bzl create mode 100644 test/unit/lto/BUILD.bazel create mode 100644 test/unit/lto/lto_test_suite.bzl diff --git a/rust/private/lto.bzl b/rust/private/lto.bzl new file mode 100644 index 0000000000..3f5c16264a --- /dev/null +++ b/rust/private/lto.bzl @@ -0,0 +1,114 @@ +"""A module defining Rust link time optimization (lto) rules""" + +load("//rust/private:utils.bzl", "is_exec_configuration") + +_LTO_MODES = [ + # Default. No mode has been explicitly set, rustc will do "thin local" LTO + # between the codegen units of a single crate. + "unspecified", + # LTO has been explicitly turned "off". + "off", + # Perform "thin" LTO. This is similar to "fat" but takes significantly less + # time to run, but provides similar performance improvements. + # + # See: + "thin", + # Perform "fat"/full LTO. + "fat", +] + +RustLtoInfo = provider( + doc = "A provider describing the link time optimization setting.", + fields = {"mode": "string: The LTO mode specified via a build setting."}, +) + +def _rust_lto_flag_impl(ctx): + value = ctx.build_setting_value + + if value not in _LTO_MODES: + msg = "{NAME} build setting allowed to take values [{EXPECTED}], but was set to: {ACTUAL}".format( + NAME = ctx.label, + VALUES = ", ".join(["'{}'".format(m) for m in _LTO_MODES]), + ACTUAL = value, + ) + fail(msg) + + return RustLtoInfo(mode = value) + +rust_lto_flag = rule( + doc = "A build setting which specifies the link time optimization mode used when building Rust code. Allowed values are: ".format(_LTO_MODES), + implementation = _rust_lto_flag_impl, + build_setting = config.string(flag = True), +) + +def _determine_lto_object_format(ctx, toolchain, crate_info): + """Determines if we should run LTO and what bitcode should get included in a built artifact. + + Args: + ctx (ctx): The calling rule's context object. + toolchain (rust_toolchain): The current target's `rust_toolchain`. + crate_info (CrateInfo): The CrateInfo provider of the target crate. + + Returns: + string: Returns one of only_object, only_bitcode, object_and_bitcode. + """ + + # Even if LTO is enabled don't use it for actions being built in the exec + # configuration, e.g. build scripts and proc-macros. This mimics Cargo. + if is_exec_configuration(ctx): + return "only_object" + + mode = toolchain._lto.mode + + if mode in ["off", "unspecified"]: + return "only_object" + + perform_linking = crate_info.type in ["bin", "staticlib", "cdylib"] + + # is_linkable = crate_info.type in ["lib", "rlib", "dylib", "proc-macro"] + is_dynamic = crate_info.type in ["dylib", "cdylib", "proc-macro"] + needs_object = perform_linking or is_dynamic + + # At this point we know LTO is enabled, otherwise we would have returned above. + + if not needs_object: + # If we're building an 'rlib' and LTO is enabled, then we can skip + # generating object files entirely. + return "only_bitcode" + elif crate_info.type == "dylib": + # If we're a dylib and we're running LTO, then only emit object code + # because 'rustc' doesn't currently support LTO with dylibs. + return "only_object" + else: + return "object_and_bitcode" + +def construct_lto_arguments(ctx, toolchain, crate_info): + """Returns a list of 'rustc' flags to configure link time optimization. + + Args: + ctx (ctx): The calling rule's context object. + toolchain (rust_toolchain): The current target's `rust_toolchain`. + crate_info (CrateInfo): The CrateInfo provider of the target crate. + + Returns: + list: A list of strings that are valid flags for 'rustc'. + """ + mode = toolchain._lto.mode + format = _determine_lto_object_format(ctx, toolchain, crate_info) + + args = [] + + if mode in ["thin", "fat", "off"] and not is_exec_configuration(ctx): + args.append("lto={}".format(mode)) + + if format in ["unspecified", "object_and_bitcode"]: + # Embedding LLVM bitcode in object files is `rustc's` default. + args.extend([]) + elif format in ["off", "only_object"]: + args.extend(["embed-bitcode=no"]) + elif format == "only_bitcode": + args.extend(["linker-plugin-lto"]) + else: + fail("unrecognized LTO object format {}".format(format)) + + return ["-C{}".format(arg) for arg in args] diff --git a/rust/private/rustc.bzl b/rust/private/rustc.bzl index 8731cc64ad..2bc6ce1959 100644 --- a/rust/private/rustc.bzl +++ b/rust/private/rustc.bzl @@ -24,6 +24,7 @@ load( ) load("//rust/private:common.bzl", "rust_common") load("//rust/private:compat.bzl", "abs") +load("//rust/private:lto.bzl", "construct_lto_arguments") load("//rust/private:providers.bzl", "RustcOutputDiagnosticsInfo", _BuildInfo = "BuildInfo") load("//rust/private:stamp.bzl", "is_stamping_enabled") load( @@ -998,6 +999,7 @@ def construct_arguments( data_paths = depset(direct = getattr(attr, "data", []), transitive = [crate_info.compile_data_targets]).to_list() add_edition_flags(rustc_flags, crate_info) + _add_lto_flags(ctx, toolchain, rustc_flags, crate_info) # Link! if ("link" in emit and crate_info.type not in ["rlib", "lib"]) or add_flags_for_binary: @@ -1583,6 +1585,18 @@ def _collect_nonstatic_linker_inputs(cc_info): )) return shared_linker_inputs +def _add_lto_flags(ctx, toolchain, args, crate): + """Adds flags to an Args object to configure LTO for 'rustc'. + + Args: + ctx (ctx): The calling rule's context object. + toolchain (rust_toolchain): The current target's `rust_toolchain`. + args (Args): A reference to an Args object + crate (CrateInfo): A CrateInfo provider + """ + lto_args = construct_lto_arguments(ctx, toolchain, crate) + args.add_all(lto_args) + def establish_cc_info(ctx, attr, crate_info, toolchain, cc_toolchain, feature_configuration, interface_library): """If the produced crate is suitable yield a CcInfo to allow for interop with cc rules diff --git a/rust/settings/BUILD.bazel b/rust/settings/BUILD.bazel index c127e0f2ab..e4166e9a0e 100644 --- a/rust/settings/BUILD.bazel +++ b/rust/settings/BUILD.bazel @@ -14,6 +14,7 @@ load( "per_crate_rustc_flag", "rustc_output_diagnostics", ) +load("//rust/private:lto.bzl", "rust_lto_flag") load("//rust/private:unpretty.bzl", "rust_unpretty_flag") load(":incompatible.bzl", "incompatible_flag") @@ -48,6 +49,12 @@ rust_unpretty_flag( visibility = ["//visibility:public"], ) +rust_lto_flag( + name = "lto", + build_setting_default = "unspecified", + visibility = ["//visibility:public"], +) + # A flag controlling whether to rename first-party crates such that their names # encode the Bazel package and target name, instead of just the target name. # diff --git a/rust/toolchain.bzl b/rust/toolchain.bzl index 41b0984efc..abb93284e3 100644 --- a/rust/toolchain.bzl +++ b/rust/toolchain.bzl @@ -3,6 +3,7 @@ load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo") load("//rust/platform:triple.bzl", "triple") load("//rust/private:common.bzl", "rust_common") +load("//rust/private:lto.bzl", "RustLtoInfo") load("//rust/private:rust_analyzer.bzl", _rust_analyzer_toolchain = "rust_analyzer_toolchain") load( "//rust/private:rustfmt.bzl", @@ -517,6 +518,7 @@ def _rust_toolchain_impl(ctx): third_party_dir = ctx.attr._third_party_dir[BuildSettingInfo].value pipelined_compilation = ctx.attr._pipelined_compilation[BuildSettingInfo].value no_std = ctx.attr._no_std[BuildSettingInfo].value + lto = ctx.attr._lto[RustLtoInfo] experimental_use_global_allocator = ctx.attr._experimental_use_global_allocator[BuildSettingInfo].value if _experimental_use_cc_common_link(ctx): @@ -701,6 +703,7 @@ def _rust_toolchain_impl(ctx): _toolchain_generated_sysroot = ctx.attr._toolchain_generated_sysroot[BuildSettingInfo].value, _incompatible_do_not_include_data_in_compile_data = ctx.attr._incompatible_do_not_include_data_in_compile_data[IncompatibleFlagInfo].enabled, _no_std = no_std, + _lto = lto, ) return [ toolchain, @@ -891,6 +894,10 @@ rust_toolchain = rule( default = Label("//rust/settings:incompatible_do_not_include_data_in_compile_data"), doc = "Label to a boolean build setting that controls whether to include data files in compile_data.", ), + "_lto": attr.label( + providers = [RustLtoInfo], + default = Label("//rust/settings:lto"), + ), "_no_std": attr.label( default = Label("//rust/settings:no_std"), ), diff --git a/test/unit/lto/BUILD.bazel b/test/unit/lto/BUILD.bazel new file mode 100644 index 0000000000..228e7840e4 --- /dev/null +++ b/test/unit/lto/BUILD.bazel @@ -0,0 +1,5 @@ +load(":lto_test_suite.bzl", "lto_test_suite") + +lto_test_suite( + name = "lto_test_suite", +) diff --git a/test/unit/lto/lto_test_suite.bzl b/test/unit/lto/lto_test_suite.bzl new file mode 100644 index 0000000000..d70d3bce24 --- /dev/null +++ b/test/unit/lto/lto_test_suite.bzl @@ -0,0 +1,123 @@ +"""Starlark tests for `//rust/settings/lto`""" + +load("@bazel_skylib//lib:unittest.bzl", "analysistest") +load("@bazel_skylib//rules:write_file.bzl", "write_file") +load("//rust:defs.bzl", "rust_library") +load( + "//test/unit:common.bzl", + "assert_action_mnemonic", + "assert_argv_contains", + "assert_argv_contains_not", + "assert_argv_contains_prefix_not", +) + +def _lto_test_impl(ctx, lto_setting, embed_bitcode, linker_plugin): + env = analysistest.begin(ctx) + target = analysistest.target_under_test(env) + + action = target.actions[0] + assert_action_mnemonic(env, action, "Rustc") + + # Check if LTO is enabled. + if lto_setting: + assert_argv_contains(env, action, "-Clto={}".format(lto_setting)) + else: + assert_argv_contains_prefix_not(env, action, "-Clto") + + # Check if we should embed bitcode. + if embed_bitcode: + assert_argv_contains(env, action, "-Cembed-bitcode={}".format(embed_bitcode)) + else: + assert_argv_contains_prefix_not(env, action, "-Cembed-bitcode") + + # Check if we should use linker plugin LTO. + if linker_plugin: + assert_argv_contains(env, action, "-Clinker-plugin-lto") + else: + assert_argv_contains_not(env, action, "-Clinker-plugin-lto") + + return analysistest.end(env) + +def _lto_level_default(ctx): + return _lto_test_impl(ctx, None, "no", False) + +_lto_level_default_test = analysistest.make( + _lto_level_default, + config_settings = {}, +) + +def _lto_level_off(ctx): + return _lto_test_impl(ctx, "off", "no", False) + +_lto_level_off_test = analysistest.make( + _lto_level_off, + config_settings = {str(Label("//rust/settings:lto")): "off"}, +) + +def _lto_level_thin(ctx): + return _lto_test_impl(ctx, "thin", None, True) + +_lto_level_thin_test = analysistest.make( + _lto_level_thin, + config_settings = {str(Label("//rust/settings:lto")): "thin"}, +) + +def _lto_level_fat(ctx): + return _lto_test_impl(ctx, "fat", None, True) + +_lto_level_fat_test = analysistest.make( + _lto_level_fat, + config_settings = {str(Label("//rust/settings:lto")): "fat"}, +) + +def lto_test_suite(name): + """Entry-point macro called from the BUILD file. + + Args: + name (str): The name of the test suite. + """ + write_file( + name = "crate_lib", + out = "lib.rs", + content = [ + "#[allow(dead_code)]", + "fn add() {}", + "", + ], + ) + + rust_library( + name = "lib", + srcs = [":lib.rs"], + edition = "2021", + ) + + _lto_level_default_test( + name = "lto_level_default_test", + target_under_test = ":lib", + ) + + _lto_level_off_test( + name = "lto_level_off_test", + target_under_test = ":lib", + ) + + _lto_level_thin_test( + name = "lto_level_thin_test", + target_under_test = ":lib", + ) + + _lto_level_fat_test( + name = "lto_level_fat_test", + target_under_test = ":lib", + ) + + native.test_suite( + name = name, + tests = [ + ":lto_level_default_test", + ":lto_level_off_test", + ":lto_level_thin_test", + ":lto_level_fat_test", + ], + )