Added codegen_units flag for setting Rustc resource sets

Author: UebelAndre

rust/private/rustc.bzl

@@ -22,13 +22,14 @@ load(
     "CPP_LINK_NODEPS_DYNAMIC_LIBRARY_ACTION_NAME",
     "CPP_LINK_STATIC_LIBRARY_ACTION_NAME",
 )
-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(":common.bzl", "rust_common")
+load(":compat.bzl", "abs")
+load(":lto.bzl", "construct_lto_arguments")
+load(":providers.bzl", "RustcOutputDiagnosticsInfo", _BuildInfo = "BuildInfo")
+load(":rustc_resource_set.bzl", "get_rustc_resource_set", "is_codegen_units_enabled")
+load(":stamp.bzl", "is_stamping_enabled")
 load(
-    "//rust/private:utils.bzl",
+    ":utils.bzl",
     "expand_dict_value_locations",
     "expand_list_element_locations",
     "find_cc_toolchain",
@@ -792,6 +793,7 @@ def collect_inputs(
     return compile_inputs, out_dir, build_env_files, build_flags_files, linkstamp_outs, ambiguous_libs
 
 def construct_arguments(
+        *,
         ctx,
         attr,
         file,
@@ -1000,6 +1002,7 @@ def construct_arguments(
 
     add_edition_flags(rustc_flags, crate_info)
     _add_lto_flags(ctx, toolchain, rustc_flags, crate_info)
+    _add_codegen_units_flags(toolchain, rustc_flags)
 
     # Link!
     if ("link" in emit and crate_info.type not in ["rlib", "lib"]) or add_flags_for_binary:
@@ -1124,6 +1127,7 @@ def construct_arguments(
     return args, env
 
 def rustc_compile_action(
+        *,
         ctx,
         attr,
         toolchain,
@@ -1337,6 +1341,7 @@ def rustc_compile_action(
                 len(crate_info.srcs.to_list()),
             ),
             toolchain = "@rules_rust//rust:toolchain_type",
+            resource_set = get_rustc_resource_set(toolchain),
         )
         if args_metadata:
             ctx.actions.run(
@@ -1372,6 +1377,7 @@ def rustc_compile_action(
                 len(crate_info.srcs.to_list()),
             ),
             toolchain = "@rules_rust//rust:toolchain_type",
+            resource_set = get_rustc_resource_set(toolchain),
         )
     else:
         fail("No process wrapper was defined for {}".format(ctx.label))
@@ -1597,6 +1603,20 @@ def _add_lto_flags(ctx, toolchain, args, crate):
     lto_args = construct_lto_arguments(ctx, toolchain, crate)
     args.add_all(lto_args)
 
+def _add_codegen_units_flags(toolchain, args):
+    """Adds flags to an Args object to configure codgen_units for 'rustc'.
+
+    https://doc.rust-lang.org/rustc/codegen-options/index.html#codegen-units
+
+    Args:
+        toolchain (rust_toolchain): The current target's `rust_toolchain`.
+        args (Args): A reference to an Args object
+    """
+    if not is_codegen_units_enabled(toolchain):
+        return
+
+    args.add("-Ccodegen-units={}".format(toolchain._codegen_units))
+
 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
 

rust/private/rustc_resource_set.bzl

@@ -0,0 +1,296 @@
+"""Resource set definitions for Rustc actions"""
+
+def _resource_set_cpu_1(_os_name, _inputs):
+    return {"cpu": 1}
+
+def _resource_set_cpu_2(_os_name, _inputs):
+    return {"cpu": 2}
+
+def _resource_set_cpu_3(_os_name, _inputs):
+    return {"cpu": 3}
+
+def _resource_set_cpu_4(_os_name, _inputs):
+    return {"cpu": 4}
+
+def _resource_set_cpu_5(_os_name, _inputs):
+    return {"cpu": 5}
+
+def _resource_set_cpu_6(_os_name, _inputs):
+    return {"cpu": 6}
+
+def _resource_set_cpu_7(_os_name, _inputs):
+    return {"cpu": 7}
+
+def _resource_set_cpu_8(_os_name, _inputs):
+    return {"cpu": 8}
+
+def _resource_set_cpu_9(_os_name, _inputs):
+    return {"cpu": 9}
+
+def _resource_set_cpu_10(_os_name, _inputs):
+    return {"cpu": 10}
+
+def _resource_set_cpu_11(_os_name, _inputs):
+    return {"cpu": 11}
+
+def _resource_set_cpu_12(_os_name, _inputs):
+    return {"cpu": 12}
+
+def _resource_set_cpu_13(_os_name, _inputs):
+    return {"cpu": 13}
+
+def _resource_set_cpu_14(_os_name, _inputs):
+    return {"cpu": 14}
+
+def _resource_set_cpu_15(_os_name, _inputs):
+    return {"cpu": 15}
+
+def _resource_set_cpu_16(_os_name, _inputs):
+    return {"cpu": 16}
+
+def _resource_set_cpu_17(_os_name, _inputs):
+    return {"cpu": 17}
+
+def _resource_set_cpu_18(_os_name, _inputs):
+    return {"cpu": 18}
+
+def _resource_set_cpu_19(_os_name, _inputs):
+    return {"cpu": 19}
+
+def _resource_set_cpu_20(_os_name, _inputs):
+    return {"cpu": 20}
+
+def _resource_set_cpu_21(_os_name, _inputs):
+    return {"cpu": 21}
+
+def _resource_set_cpu_22(_os_name, _inputs):
+    return {"cpu": 22}
+
+def _resource_set_cpu_23(_os_name, _inputs):
+    return {"cpu": 23}
+
+def _resource_set_cpu_24(_os_name, _inputs):
+    return {"cpu": 24}
+
+def _resource_set_cpu_25(_os_name, _inputs):
+    return {"cpu": 25}
+
+def _resource_set_cpu_26(_os_name, _inputs):
+    return {"cpu": 26}
+
+def _resource_set_cpu_27(_os_name, _inputs):
+    return {"cpu": 27}
+
+def _resource_set_cpu_28(_os_name, _inputs):
+    return {"cpu": 28}
+
+def _resource_set_cpu_29(_os_name, _inputs):
+    return {"cpu": 29}
+
+def _resource_set_cpu_30(_os_name, _inputs):
+    return {"cpu": 30}
+
+def _resource_set_cpu_31(_os_name, _inputs):
+    return {"cpu": 31}
+
+def _resource_set_cpu_32(_os_name, _inputs):
+    return {"cpu": 32}
+
+def _resource_set_cpu_33(_os_name, _inputs):
+    return {"cpu": 33}
+
+def _resource_set_cpu_34(_os_name, _inputs):
+    return {"cpu": 34}
+
+def _resource_set_cpu_35(_os_name, _inputs):
+    return {"cpu": 35}
+
+def _resource_set_cpu_36(_os_name, _inputs):
+    return {"cpu": 36}
+
+def _resource_set_cpu_37(_os_name, _inputs):
+    return {"cpu": 37}
+
+def _resource_set_cpu_38(_os_name, _inputs):
+    return {"cpu": 38}
+
+def _resource_set_cpu_39(_os_name, _inputs):
+    return {"cpu": 39}
+
+def _resource_set_cpu_40(_os_name, _inputs):
+    return {"cpu": 40}
+
+def _resource_set_cpu_41(_os_name, _inputs):
+    return {"cpu": 41}
+
+def _resource_set_cpu_42(_os_name, _inputs):
+    return {"cpu": 42}
+
+def _resource_set_cpu_43(_os_name, _inputs):
+    return {"cpu": 43}
+
+def _resource_set_cpu_44(_os_name, _inputs):
+    return {"cpu": 44}
+
+def _resource_set_cpu_45(_os_name, _inputs):
+    return {"cpu": 45}
+
+def _resource_set_cpu_46(_os_name, _inputs):
+    return {"cpu": 46}
+
+def _resource_set_cpu_47(_os_name, _inputs):
+    return {"cpu": 47}
+
+def _resource_set_cpu_48(_os_name, _inputs):
+    return {"cpu": 48}
+
+def _resource_set_cpu_49(_os_name, _inputs):
+    return {"cpu": 49}
+
+def _resource_set_cpu_50(_os_name, _inputs):
+    return {"cpu": 50}
+
+def _resource_set_cpu_51(_os_name, _inputs):
+    return {"cpu": 51}
+
+def _resource_set_cpu_52(_os_name, _inputs):
+    return {"cpu": 52}
+
+def _resource_set_cpu_53(_os_name, _inputs):
+    return {"cpu": 53}
+
+def _resource_set_cpu_54(_os_name, _inputs):
+    return {"cpu": 54}
+
+def _resource_set_cpu_55(_os_name, _inputs):
+    return {"cpu": 55}
+
+def _resource_set_cpu_56(_os_name, _inputs):
+    return {"cpu": 56}
+
+def _resource_set_cpu_57(_os_name, _inputs):
+    return {"cpu": 57}
+
+def _resource_set_cpu_58(_os_name, _inputs):
+    return {"cpu": 58}
+
+def _resource_set_cpu_59(_os_name, _inputs):
+    return {"cpu": 59}
+
+def _resource_set_cpu_60(_os_name, _inputs):
+    return {"cpu": 60}
+
+def _resource_set_cpu_61(_os_name, _inputs):
+    return {"cpu": 61}
+
+def _resource_set_cpu_62(_os_name, _inputs):
+    return {"cpu": 62}
+
+def _resource_set_cpu_63(_os_name, _inputs):
+    return {"cpu": 63}
+
+def _resource_set_cpu_64(_os_name, _inputs):
+    return {"cpu": 64}
+
+_RESOURCE_SETS = {
+    1: _resource_set_cpu_1,
+    2: _resource_set_cpu_2,
+    3: _resource_set_cpu_3,
+    4: _resource_set_cpu_4,
+    5: _resource_set_cpu_5,
+    6: _resource_set_cpu_6,
+    7: _resource_set_cpu_7,
+    8: _resource_set_cpu_8,
+    9: _resource_set_cpu_9,
+    10: _resource_set_cpu_10,
+    11: _resource_set_cpu_11,
+    12: _resource_set_cpu_12,
+    13: _resource_set_cpu_13,
+    14: _resource_set_cpu_14,
+    15: _resource_set_cpu_15,
+    16: _resource_set_cpu_16,
+    17: _resource_set_cpu_17,
+    18: _resource_set_cpu_18,
+    19: _resource_set_cpu_19,
+    20: _resource_set_cpu_20,
+    21: _resource_set_cpu_21,
+    22: _resource_set_cpu_22,
+    23: _resource_set_cpu_23,
+    24: _resource_set_cpu_24,
+    25: _resource_set_cpu_25,
+    26: _resource_set_cpu_26,
+    27: _resource_set_cpu_27,
+    28: _resource_set_cpu_28,
+    29: _resource_set_cpu_29,
+    30: _resource_set_cpu_30,
+    31: _resource_set_cpu_31,
+    32: _resource_set_cpu_32,
+    33: _resource_set_cpu_33,
+    34: _resource_set_cpu_34,
+    35: _resource_set_cpu_35,
+    36: _resource_set_cpu_36,
+    37: _resource_set_cpu_37,
+    38: _resource_set_cpu_38,
+    39: _resource_set_cpu_39,
+    40: _resource_set_cpu_40,
+    41: _resource_set_cpu_41,
+    42: _resource_set_cpu_42,
+    43: _resource_set_cpu_43,
+    44: _resource_set_cpu_44,
+    45: _resource_set_cpu_45,
+    46: _resource_set_cpu_46,
+    47: _resource_set_cpu_47,
+    48: _resource_set_cpu_48,
+    49: _resource_set_cpu_49,
+    50: _resource_set_cpu_50,
+    51: _resource_set_cpu_51,
+    52: _resource_set_cpu_52,
+    53: _resource_set_cpu_53,
+    54: _resource_set_cpu_54,
+    55: _resource_set_cpu_55,
+    56: _resource_set_cpu_56,
+    57: _resource_set_cpu_57,
+    58: _resource_set_cpu_58,
+    59: _resource_set_cpu_59,
+    60: _resource_set_cpu_60,
+    61: _resource_set_cpu_61,
+    62: _resource_set_cpu_62,
+    63: _resource_set_cpu_63,
+    64: _resource_set_cpu_64,
+}
+
+def is_codegen_units_enabled(toolchain):
+    """Check whether or not codegen-units should be applied by the toolchain.
+
+    Args:
+        toolchain (rust_toolchain): The current `rust_toolchain`.
+
+    Returns:
+        bool: Whether or not codegen-units is enabled
+    """
+    if toolchain._codegen_units <= 0:
+        return False
+
+    if toolchain._experimental_use_cc_common_link:
+        return False
+
+    return True
+
+def get_rustc_resource_set(toolchain):
+    """Get the `ctx.actions.run.resource_set` for the `Rustc` action.
+
+    Args:
+        toolchain (rust_toolchain): The current rust_toolchain toolchain.
+
+    Returns:
+        Optional[Callable]: A resource set appropriate for the current configuration.
+    """
+    if not is_codegen_units_enabled(toolchain):
+        return None
+
+    codegen_units = toolchain._codegen_units
+
+    if codegen_units > len(_RESOURCE_SETS):
+        return _RESOURCE_SETS[len(_RESOURCE_SETS)]
+
+    return _RESOURCE_SETS[codegen_units]