Add "proto_profile" attribute to fdo_profile rule.

PiperOrigin-RevId: 212653098

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

@@ -332,6 +332,8 @@ public ConfiguredTarget create(RuleContext ruleContext)
 
     PathFragment fdoZip = null;
     FdoInputFile prefetchHints = null;
+    Artifact protoProfileArtifact = null;
+    boolean allowInference = true;
     if (configuration.getCompilationMode() == CompilationMode.OPT) {
       if (cppConfiguration.getFdoPrefetchHintsLabel() != null) {
         FdoPrefetchHintsProvider provider =
@@ -342,32 +344,50 @@ public ConfiguredTarget create(RuleContext ruleContext)
       if (cppConfiguration.getFdoPath() != null) {
         fdoZip = cppConfiguration.getFdoPath();
       } else if (cppConfiguration.getFdoOptimizeLabel() != null) {
-        Artifact fdoArtifact =
-            ruleContext.getPrerequisiteArtifact(CcToolchainRule.FDO_OPTIMIZE_ATTR, Mode.TARGET);
-        if (!fdoArtifact.isSourceArtifact()) {
-          ruleContext.ruleError("--fdo_optimize points to a target that is not an input file");
-          return null;
-        }
-        Label fdoLabel =
-            ruleContext.getPrerequisite(CcToolchainRule.FDO_OPTIMIZE_ATTR, Mode.TARGET).getLabel();
-        if (!fdoLabel
-            .getPackageIdentifier()
-            .getPathUnderExecRoot()
-            .getRelative(fdoLabel.getName())
-            .equals(fdoArtifact.getExecPath())) {
-          ruleContext.ruleError("--fdo_optimize points to a target that is not an input file");
-          return null;
+        // If fdo_profile rule is used, do not allow inferring proto.profile from AFDO profile.
+        allowInference = false;
+
+        FdoProfileProvider fdoProfileProvider =
+            ruleContext.getPrerequisite(
+                CcToolchainRule.FDO_OPTIMIZE_ATTR, Mode.TARGET, FdoProfileProvider.PROVIDER);
+        if (fdoProfileProvider != null) {
+          fdoZip = fdoProfileProvider.getProfilePathFragment();
+          protoProfileArtifact = fdoProfileProvider.getProtoProfileArtifact();
+        } else {
+          ImmutableList<Artifact> fdoArtifacts =
+              ruleContext
+                  .getPrerequisiteArtifacts(CcToolchainRule.FDO_OPTIMIZE_ATTR, Mode.TARGET)
+                  .list();
+          if (fdoArtifacts.size() != 1) {
+            ruleContext.ruleError("--fdo_optimize does not point to a single target");
+            return null;
+          }
+
+          Artifact fdoArtifact = fdoArtifacts.get(0);
+          if (!fdoArtifact.isSourceArtifact()) {
+            ruleContext.ruleError("--fdo_optimize points to a target that is not an input file");
+            return null;
+          }
+          Label fdoLabel =
+              ruleContext
+                  .getPrerequisite(CcToolchainRule.FDO_OPTIMIZE_ATTR, Mode.TARGET)
+                  .getLabel();
+          if (!fdoLabel
+              .getPackageIdentifier()
+              .getPathUnderExecRoot()
+              .getRelative(fdoLabel.getName())
+              .equals(fdoArtifact.getExecPath())) {
+            ruleContext.ruleError("--fdo_optimize points to a target that is not an input file");
+            return null;
+          }
+          fdoZip = fdoArtifact.getPath().asFragment();
         }
-        fdoZip = fdoArtifact.getPath().asFragment();
       } else if (cppConfiguration.getFdoProfileLabel() != null) {
         FdoProfileProvider fdoProvider =
             ruleContext.getPrerequisite(
                 CcToolchainRule.FDO_PROFILE_ATTR, Mode.TARGET, FdoProfileProvider.PROVIDER);
-        FdoInputFile inputFile = fdoProvider.getInputFile();
-        fdoZip =
-            inputFile.getAbsolutePath() != null
-                ? inputFile.getAbsolutePath()
-                : inputFile.getArtifact().getPath().asFragment();
+        fdoZip = fdoProvider.getProfilePathFragment();
+        protoProfileArtifact = fdoProvider.getProtoProfileArtifact();
       }
     }
 
@@ -594,7 +614,9 @@ public ConfiguredTarget create(RuleContext ruleContext)
                 fdoMode,
                 cppConfiguration.getFdoInstrument(),
                 profileArtifact,
-                prefetchHintsArtifact),
+                prefetchHintsArtifact,
+                protoProfileArtifact,
+                allowInference),
             cppConfiguration.useLLVMCoverageMapFormat(),
             configuration.isCodeCoverageEnabled(),
             configuration.isHostConfiguration());

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

@@ -180,7 +180,7 @@ public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment env)
         // TODO(b/78578234): Make this the default and remove the late-bound versions.
         .add(attr("libc_top", LABEL).allowedFileTypes())
         .add(attr(LIBC_TOP_ATTR, LABEL).value(LIBC_TOP_VALUE))
-        .add(attr(FDO_OPTIMIZE_ATTR, LABEL).singleArtifact().value(FDO_OPTIMIZE_VALUE))
+        .add(attr(FDO_OPTIMIZE_ATTR, LABEL).value(FDO_OPTIMIZE_VALUE))
         .add(
             attr(FDO_PROFILE_ATTR, LABEL)
                 .allowedRuleClasses("fdo_profile")

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

@@ -13,13 +13,15 @@
 // limitations under the License.
 package com.google.devtools.build.lib.rules.cpp;
 
+import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.MutableActionGraph.ActionConflictException;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.RuleConfiguredTargetBuilder;
 import com.google.devtools.build.lib.analysis.RuleConfiguredTargetFactory;
 import com.google.devtools.build.lib.analysis.RuleContext;
 import com.google.devtools.build.lib.analysis.Runfiles;
 import com.google.devtools.build.lib.analysis.RunfilesProvider;
+import com.google.devtools.build.lib.analysis.configuredtargets.RuleConfiguredTarget.Mode;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 
 /** Implementation for the {@code fdo_profile} rule. */
@@ -34,8 +36,14 @@ public ConfiguredTarget create(RuleContext ruleContext)
       return null;
     }
 
+    Artifact protoProfileArtifact =
+        ruleContext.getPrerequisiteArtifact("proto_profile", Mode.TARGET);
+    if (protoProfileArtifact != null && !protoProfileArtifact.isSourceArtifact()) {
+      ruleContext.attributeError("proto_profile", "the target is not an input file");
+    }
+
     return new RuleConfiguredTargetBuilder(ruleContext)
-        .addNativeDeclaredProvider(new FdoProfileProvider(inputFile))
+        .addNativeDeclaredProvider(new FdoProfileProvider(inputFile, protoProfileArtifact))
         .addProvider(RunfilesProvider.simple(Runfiles.EMPTY))
         .build();
   }

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

@@ -13,9 +13,11 @@
 // limitations under the License.
 package com.google.devtools.build.lib.rules.cpp;
 
+import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.packages.NativeInfo;
 import com.google.devtools.build.lib.packages.NativeProvider;
+import com.google.devtools.build.lib.vfs.PathFragment;
 
 /** Provider that contains the profile used for FDO. */
 @Immutable
@@ -25,13 +27,25 @@ public final class FdoProfileProvider extends NativeInfo {
       new NativeProvider<FdoProfileProvider>(FdoProfileProvider.class, "FdoProfileInfo") {};
 
   private final FdoInputFile fdoInputFile;
+  private final Artifact protoProfileArtifact;
 
-  public FdoProfileProvider(FdoInputFile fdoInputFile) {
+  public FdoProfileProvider(FdoInputFile fdoInputFile, Artifact protoProfileArtifact) {
     super(PROVIDER);
     this.fdoInputFile = fdoInputFile;
+    this.protoProfileArtifact = protoProfileArtifact;
   }
 
   public FdoInputFile getInputFile() {
     return fdoInputFile;
   }
+
+  public PathFragment getProfilePathFragment() {
+    return fdoInputFile.getAbsolutePath() != null
+        ? fdoInputFile.getAbsolutePath()
+        : fdoInputFile.getArtifact().getPath().asFragment();
+  }
+
+  public Artifact getProtoProfileArtifact() {
+    return protoProfileArtifact;
+  }
 }

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

@@ -51,6 +51,12 @@ public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment env)
         that holds an LLVM profraw profile, or .afdo for AutoFDO profile.
         <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
         .add(attr("absolute_path_profile", Type.STRING))
+        /* <!-- #BLAZE_RULE(fdo_profile).ATTRIBUTE(proto_profile) -->
+        Label of the protobuf profile.
+        <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
+        .add(attr("proto_profile", LABEL)
+            .allowedFileTypes(FileTypeSet.ANY_FILE)
+            .singleArtifact())
         .advertiseProvider(FdoProfileProvider.class)
         .build();
   }

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

@@ -57,15 +57,27 @@ public enum FdoMode {
   private final String fdoInstrument;
   private final Artifact profileArtifact;
   private final Artifact prefetchHintsArtifact;
+  private final Artifact protoProfileArtifact;
+  // If true, proto.profile location is inferred from AFDO profile.
+  // TODO(seongkim): remove it after transtition to fdo_profile rule completes.
+  private final boolean allowInference;
 
   @AutoCodec.Instantiator
-  public FdoProvider(Path fdoProfilePath, FdoMode fdoMode, String fdoInstrument,
-      Artifact profileArtifact, Artifact prefetchHintsArtifact) {
+  public FdoProvider(
+      Path fdoProfilePath,
+      FdoMode fdoMode,
+      String fdoInstrument,
+      Artifact profileArtifact,
+      Artifact prefetchHintsArtifact,
+      Artifact protoProfileArtifact,
+      boolean allowInference) {
     this.fdoProfilePath = fdoProfilePath;
     this.fdoMode = fdoMode;
     this.fdoInstrument = fdoInstrument;
     this.profileArtifact = profileArtifact;
     this.prefetchHintsArtifact = prefetchHintsArtifact;
+    this.protoProfileArtifact = protoProfileArtifact;
+    this.allowInference = allowInference;
   }
 
   /** <b>DO NOT EVER USE PATHS IN THE ANALYSIS PHASE!
@@ -95,4 +107,12 @@ public Artifact getProfileArtifact() {
   public Artifact getPrefetchHintsArtifact() {
     return prefetchHintsArtifact;
   }
+
+  public Artifact getProtoProfileArtifact() {
+    return protoProfileArtifact;
+  }
+
+  public boolean allowInference() {
+    return allowInference;
+  }
 }