Support passing LLVM MemProf (memory profile) to CppCompile

Adds support for supplying a MemProf profile, which can be supplied to LLVM to
automatically hint allocations based on profiled memory behavior. The patch
adds Blaze flag --memprof_profile=<label>, along with a memprof_profile rule
that allows specifying a profile either by absolute path or label.

The profile can either be an LLVM memprof profile with a .profdata extension,
or a zipfile containing one. The profile is symlinked and optionally unzipped,
and set in the memprof_profile_path build variable, which can be used to pass
the profile to the compiler via a feature.

PiperOrigin-RevId: 546311130
Change-Id: I896898bf3bddcf42214cde3f3f27d2ed18de6b26

Author: Googler

src/main/java/com/google/devtools/build/lib/bazel/rules/CcRules.java

@@ -43,6 +43,7 @@
 import com.google.devtools.build.lib.rules.cpp.DebugPackageProvider;
 import com.google.devtools.build.lib.rules.cpp.FdoPrefetchHintsRule;
 import com.google.devtools.build.lib.rules.cpp.FdoProfileRule;
+import com.google.devtools.build.lib.rules.cpp.MemProfProfileRule;
 import com.google.devtools.build.lib.rules.cpp.PropellerOptimizeRule;
 import com.google.devtools.build.lib.rules.platform.PlatformRules;
 import com.google.devtools.build.lib.starlarkbuildapi.cpp.CcBootstrap;
@@ -95,6 +96,7 @@ public void init(ConfiguredRuleClassProvider.Builder builder) {
     builder.addRuleDefinition(new FdoProfileRule());
     builder.addRuleDefinition(new FdoPrefetchHintsRule());
     builder.addRuleDefinition(new CcLinkingRule());
+    builder.addRuleDefinition(new MemProfProfileRule());
     builder.addRuleDefinition(new PropellerOptimizeRule());
     builder.addStarlarkBuiltinsInternal(
         "StaticallyLinkedMarkerProvider", StaticallyLinkedMarkerProvider.PROVIDER);

src/main/java/com/google/devtools/build/lib/rules/cpp/CcCommon.java

@@ -771,6 +771,10 @@ public static FeatureConfiguration configureFeaturesOrThrowEvalException(
       allRequestedFeaturesBuilder.add(CppRuleClasses.PROPELLER_OPTIMIZE);
     }
 
+    if (cppConfiguration.getMemProfProfileLabel() != null) {
+      allRequestedFeaturesBuilder.add(CppRuleClasses.MEMPROF_OPTIMIZE);
+    }
+
     for (String feature : allFeatures.build()) {
       if (!allUnsupportedFeatures.contains(feature)) {
         allRequestedFeaturesBuilder.add(feature);

src/main/java/com/google/devtools/build/lib/rules/cpp/CcCompilationHelper.java

@@ -120,6 +120,12 @@ private void configureFdoBuildVariables(
       }
     }
 
+    if (fdoContext.getMemProfProfileArtifact() != null) {
+      variablesBuilder.put(
+          CompileBuildVariables.MEMPROF_PROFILE_PATH.getVariableName(),
+          fdoContext.getMemProfProfileArtifact().getExecPathString());
+    }
+
     FdoContext.BranchFdoProfile branchFdoProfile = fdoContext.getBranchFdoProfile();
     // Optimization phase
     if (branchFdoProfile != null) {
@@ -173,6 +179,9 @@ private NestedSet<Artifact> getAuxiliaryFdoInputs() {
         auxiliaryInputs.add(fdoContext.getPropellerOptimizeInputFile().getLdArtifact());
       }
     }
+    if (fdoContext.getMemProfProfileArtifact() != null) {
+      auxiliaryInputs.add(fdoContext.getMemProfProfileArtifact());
+    }
     FdoContext.BranchFdoProfile branchFdoProfile = fdoContext.getBranchFdoProfile();
     // If --fdo_optimize was not specified, we don't have any additional inputs.
     if (branchFdoProfile != null) {

src/main/java/com/google/devtools/build/lib/rules/cpp/CcToolchainAttributesProvider.java

@@ -84,9 +84,11 @@ public class CcToolchainAttributesProvider extends NativeInfo implements HasCcTo
   private final ImmutableList<Artifact> fdoOptimizeArtifacts;
   private final FdoPrefetchHintsProvider fdoPrefetch;
   private final PropellerOptimizeProvider propellerOptimize;
+  private final MemProfProfileProvider memprofProfileProvider;
   private final TransitiveInfoCollection moduleMap;
   private final Artifact moduleMapArtifact;
   private final Artifact zipper;
+  private final Artifact defaultZipper;
   private final String purposePrefix;
   private final String runtimeSolibDirBase;
   private final LicensesProvider licensesProvider;
@@ -177,9 +179,13 @@ public CcToolchainAttributesProvider(
         ruleContext.getPrerequisite(":fdo_prefetch_hints", FdoPrefetchHintsProvider.PROVIDER);
     this.propellerOptimize =
         ruleContext.getPrerequisite(":propeller_optimize", PropellerOptimizeProvider.PROVIDER);
+    this.memprofProfileProvider =
+        ruleContext.getPrerequisite(
+            CcToolchainRule.MEMPROF_PROFILE_ATTR, MemProfProfileProvider.PROVIDER);
     this.moduleMap = ruleContext.getPrerequisite("module_map");
     this.moduleMapArtifact = ruleContext.getPrerequisiteArtifact("module_map");
     this.zipper = ruleContext.getPrerequisiteArtifact(":zipper");
+    this.defaultZipper = ruleContext.getPrerequisiteArtifact(":default_zipper");
     this.purposePrefix = Actions.escapeLabel(ruleContext.getLabel()) + "_";
     this.runtimeSolibDirBase = "_solib_" + "_" + Actions.escapeLabel(ruleContext.getLabel());
     this.staticRuntimeLib = ruleContext.getPrerequisite("static_runtime_lib");
@@ -274,6 +280,10 @@ public PropellerOptimizeProvider getPropellerOptimize() {
     return propellerOptimize;
   }
 
+  public MemProfProfileProvider getMemProfProfileProvider() {
+    return memprofProfileProvider;
+  }
+
   public String getToolchainIdentifier() {
     return toolchainIdentifier;
   }
@@ -432,10 +442,16 @@ public FdoProfileProvider getXFdoProfileProvider() {
     return xfdoProfileProvider;
   }
 
+  /* Get the FDO-specific zipper. */
   public Artifact getZipper() {
     return zipper;
   }
 
+  /* Get the non FDO-specific zipper. */
+  public Artifact getDefaultZipper() {
+    return defaultZipper;
+  }
+
   public NestedSet<Artifact> getFullInputsForLink() {
     return fullInputsForLink;
   }

src/main/java/com/google/devtools/build/lib/rules/cpp/CcToolchainRule.java

@@ -46,6 +46,7 @@ public final class CcToolchainRule implements RuleDefinition {
   public static final String FDO_PROFILE_ATTR = ":fdo_profile";
   public static final String CSFDO_PROFILE_ATTR = ":csfdo_profile";
   public static final String XFDO_PROFILE_ATTR = ":xfdo_profile";
+  public static final String MEMPROF_PROFILE_ATTR = ":memprof_profile";
   public static final String TOOLCHAIN_CONFIG_ATTR = "toolchain_config";
 
   private static Label getLabel(AttributeMap attributes, String attrName, Label defaultValue) {
@@ -126,6 +127,12 @@ private static Label getLabel(AttributeMap attributes, String attrName, Label de
           null,
           (rule, attributes, cppConfig) -> cppConfig.getPropellerOptimizeLabel());
 
+  private static final LabelLateBoundDefault<?> MEMPROF_PROFILE_VALUE =
+      LabelLateBoundDefault.fromTargetConfiguration(
+          CppConfiguration.class,
+          /* defaultValue= */ null,
+          (rule, attributes, cppConfig) -> cppConfig.getMemProfProfileLabel());
+
   /**
    * Returns true if zipper should be loaded. We load the zipper executable if FDO optimization is
    * enabled through --fdo_optimize or --fdo_profile
@@ -136,9 +143,14 @@ private static boolean shouldIncludeZipperInToolchain(CppConfiguration cppConfig
         || cppConfiguration.getFdoPath() != null;
   }
 
+  private static boolean shouldIncludeDefaultZipperInToolchain(CppConfiguration cppConfiguration) {
+    return cppConfiguration.getMemProfProfileLabel() != null;
+  }
+
   @Override
   public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment env) {
     final Label zipper = env.getToolsLabel("//tools/zip:unzip_fdo");
+    final Label defaultZipper = env.getToolsLabel("//tools/zip:zipper");
     return builder
         .requiresConfigurationFragments(CppConfiguration.class, PlatformConfiguration.class)
         .advertiseProvider(TemplateVariableInfo.class)
@@ -312,6 +324,9 @@ instead of having no transition (i.e. target platform by default).
         .add(
             attr(CcToolchain.CC_TOOLCHAIN_TYPE_ATTRIBUTE_NAME, NODEP_LABEL)
                 .value(CppRuleClasses.ccToolchainTypeAttribute(env)))
+        // This is the FDO-specific zipper, which takes the cpu type and extracts the raw FDO
+        // profile filename that is the closest match for that target. The name is expected
+        // to be of the format fdocontrolz_profile[-$cpu].profraw
         .add(
             attr(":zipper", LABEL)
                 .cfg(ExecutionTransitionFactory.createFactory())
@@ -322,6 +337,20 @@ instead of having no transition (i.e. target platform by default).
                         null,
                         (rule, attributes, cppConfig) ->
                             shouldIncludeZipperInToolchain(cppConfig) ? zipper : null)))
+        // This is the normal (non FDO-specific) zipper, that simply unzips the zipfile
+        // contents into the given destination directory.
+        .add(
+            attr(":default_zipper", LABEL)
+                .cfg(ExecutionTransitionFactory.createFactory())
+                .singleArtifact()
+                .value(
+                    LabelLateBoundDefault.fromTargetConfiguration(
+                        CppConfiguration.class,
+                        /* defaultValue= */ null,
+                        (rule, attributes, cppConfig) ->
+                            shouldIncludeDefaultZipperInToolchain(cppConfig)
+                                ? defaultZipper
+                                : null)))
 
         // TODO(b/78578234): Make this the default and remove the late-bound versions.
         /* <!-- #BLAZE_RULE(cc_toolchain).ATTRIBUTE(libc_top) -->
@@ -381,6 +410,12 @@ instead of having no transition (i.e. target platform by default).
                 .value(PROPELLER_OPTIMIZE)
                 // Should be in the target configuration
                 .cfg(NoTransition.createFactory()))
+        .add(
+            attr(MEMPROF_PROFILE_ATTR, LABEL)
+                .allowedRuleClasses("memprof_profile")
+                .mandatoryProviders(ImmutableList.of(MemProfProfileProvider.PROVIDER.id()))
+                .value(MEMPROF_PROFILE_VALUE)
+                .cfg(NoTransition.createFactory()))
         /* <!-- #BLAZE_RULE(cc_toolchain).ATTRIBUTE(toolchain_identifier) -->
         The identifier used to match this cc_toolchain with the corresponding
         crosstool_config.toolchain.

src/main/java/com/google/devtools/build/lib/rules/cpp/CompileBuildVariables.java

@@ -127,6 +127,8 @@ public enum CompileBuildVariables {
   PROPELLER_OPTIMIZE_CC_PATH("propeller_optimize_cc_path"),
   /** Path to the Propeller Optimize linker profile artifact */
   PROPELLER_OPTIMIZE_LD_PATH("propeller_optimize_ld_path"),
+  /** Path to the memprof profile artifact */
+  MEMPROF_PROFILE_PATH("memprof_profile_path"),
   /** Variable for includes that compiler needs to include into sources. */
   INCLUDES("includes");
 

src/main/java/com/google/devtools/build/lib/rules/cpp/CppActionConfigs.java

@@ -411,6 +411,22 @@ public static ImmutableList<CToolchain.Feature> getLegacyFeatures(
                         "  }")));
       }
 
+      if (!existingFeatureNames.contains(CppRuleClasses.MEMPROF_OPTIMIZE)) {
+        featureBuilder.add(
+            getFeature(
+                Joiner.on("\n")
+                    .join(
+                        "  name: 'memprof_optimize'",
+                        "  flag_set {",
+                        "    action: 'c-compile'",
+                        "    action: 'c++-compile'",
+                        "    flag_group {",
+                        "      expand_if_all_available: 'memprof_profile_path'",
+                        "      flag: '-memprof-profile-file=" + "%{memprof_profile_path}'",
+                        "    }",
+                        "  }")));
+      }
+
       if (!existingFeatureNames.contains(CppRuleClasses.BUILD_INTERFACE_LIBRARIES)) {
         featureBuilder.add(
             getFeature(

src/main/java/com/google/devtools/build/lib/rules/cpp/CppConfiguration.java

@@ -666,6 +666,11 @@ Label getXFdoProfileLabel() {
     return cppOptions.xfdoProfileLabel;
   }
 
+  @Nullable
+  Label getMemProfProfileLabel() {
+    return cppOptions.getMemProfProfileLabel();
+  }
+
   public boolean isFdoAbsolutePathEnabled() {
     return cppOptions.enableFdoProfileAbsolutePath;
   }

src/main/java/com/google/devtools/build/lib/rules/cpp/CppOptions.java

@@ -530,6 +530,21 @@ public Label getPropellerOptimizeLabel() {
     return propellerOptimizeLabel;
   }
 
+  @Option(
+      name = "memprof_profile",
+      defaultValue = "null",
+      converter = LabelConverter.class,
+      category = "flags",
+      documentationCategory = OptionDocumentationCategory.OUTPUT_PARAMETERS,
+      effectTags = {OptionEffectTag.AFFECTS_OUTPUTS},
+      help = "Use memprof profile.")
+  public Label memprofProfileLabel;
+
+  /** Returns the --memprof_profile value. */
+  public Label getMemProfProfileLabel() {
+    return memprofProfileLabel;
+  }
+
   @Option(
       name = "save_temps",
       defaultValue = "false",

src/main/java/com/google/devtools/build/lib/rules/cpp/CppRuleClasses.java

@@ -424,6 +424,9 @@ public static ToolchainTypeRequirement ccToolchainTypeRequirement(RuleDefinition
   /** A string constant for the propeller optimize feature. */
   public static final String PROPELLER_OPTIMIZE = "propeller_optimize";
 
+  /** A string constant for the memprof profile optimization feature. */
+  public static final String MEMPROF_OPTIMIZE = "memprof_optimize";
+
   /**
    * A string constant for the propeller_optimize_thinlto_compile_actions feature.
    *